From 9e2dfbca19d13f62f926ee6473d2640aa1fa70f4 Mon Sep 17 00:00:00 2001 From: yashagar Date: Fri, 10 Jul 2026 09:18:44 +0000 Subject: [PATCH 01/52] initial kernel and tests --- aiter/ops/flydsl/__init__.py | 8 + .../ops/flydsl/kernels/_mqa_logits_common.py | 178 ++++++ aiter/ops/flydsl/kernels/fp8_mqa_logits.py | 160 ++--- .../flydsl/kernels/fp8_paged_mqa_logits.py | 596 ++++++++++++++++++ .../test_flydsl_fp8_paged_mqa_logits.py | 304 +++++++++ .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 315 +++++++++ 6 files changed, 1438 insertions(+), 123 deletions(-) create mode 100644 aiter/ops/flydsl/kernels/_mqa_logits_common.py create mode 100644 aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py create mode 100644 op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py create mode 100644 op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 977b47c702c..26066a342bc 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -46,6 +46,11 @@ KERNEL_VARIANTS as FP8_MQA_LOGITS_VARIANTS, DEFAULT_VARIANT as FP8_MQA_LOGITS_DEFAULT_VARIANT, ) + from .kernels.fp8_paged_mqa_logits import ( + flydsl_fp8_paged_mqa_logits, + KERNEL_VARIANTS as FP8_PAGED_MQA_LOGITS_VARIANTS, + DEFAULT_VARIANT as FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT, + ) # from .linear_attention_kernels import flydsl_gdr_decode @@ -59,5 +64,8 @@ "flydsl_fp8_mqa_logits", "FP8_MQA_LOGITS_VARIANTS", "FP8_MQA_LOGITS_DEFAULT_VARIANT", + "flydsl_fp8_paged_mqa_logits", + "FP8_PAGED_MQA_LOGITS_VARIANTS", + "FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT", # "flydsl_gdr_decode", ] diff --git a/aiter/ops/flydsl/kernels/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/_mqa_logits_common.py new file mode 100644 index 00000000000..d6a49419594 --- /dev/null +++ b/aiter/ops/flydsl/kernels/_mqa_logits_common.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Shared FlyDSL building blocks for the FP8 MQA-logits kernels. + +Both the dense (``fp8_mqa_logits.py``) and paged (``fp8_paged_mqa_logits.py``) +kernels compute the same core per query row / KV column:: + + logits[row, n] = sum_h ReLU( * kv_scale[n]) * weights[row, h] + +They differ only in how K / kv_scale are addressed (contiguous vs paged gather) +and in the windowing / output layout. Everything that is identical -- the fp8 +16x16x32 MFMA compute, the ReLU * weight head-sum + kv-scale hoist, the in-wave +``shuffle_xor`` head reduce, the fp8 dword-pack load, the FN->FNUZ byte patch, +and the i64 byte-base per-row output view -- lives here so neither kernel +duplicates it. +""" + +# No `from __future__ import annotations`: FlyDSL arg typing needs real +# annotation objects, not PEP 563 strings. + +from functools import lru_cache + +import torch + +import flydsl.expr as fx +from flydsl.expr import arith, range_constexpr, rocdl +from flydsl.expr.numeric import ArithValue +from flydsl.expr.typing import T + +from .tensor_shim import GTensor, _to_raw + +Vec = fx.Vector + +# MFMA tile dims of the fp8 16x16x32 atom: MFMA_M x MFMA_N output tile, +# MFMA_K fp8 elements reduced per MFMA step. +MFMA_M = 16 +MFMA_N = 16 +MFMA_K = 32 + +# The fp8 16x16x32 MFMA intrinsic. On gfx942 it reads operands as e4m3 FNUZ +# (bias 8); on gfx950 it reads the native e4m3 FN (OCP). See the FN->FNUZ patch +# below and the per-kernel ``convert_*_fn`` gates. +MFMA_FN = rocdl.mfma_f32_16x16x32_fp8_fp8 + +DEFAULT_COMPILE_HINTS = { + "waves_per_eu": 2, + "fast_fp_math": True, +} + + +def i32_add(a, b): + """i32 add (result stays Int32, not index type).""" + return fx.Int32(arith.addi(_to_raw(a), _to_raw(b))) + + +@lru_cache(maxsize=8) +def device_cu_count(device_index: int) -> int: + """Compute-unit count for a CUDA/HIP device (cached); 304 if unavailable.""" + try: + return torch.cuda.get_device_properties(device_index).multi_processor_count + except Exception: + return 304 + + +def load_pack_i64(i32_view, byte_off_i32): + """Load 8 fp8 bytes as 2 i32 dwords and bitcast to i64. + + fp8 operands are read 8 bytes at a time as 2 i32 dwords (a ``v8i8`` + buffer_load fails to lower on gfx942), then bitcast to i64 for the MFMA. + ``i32_view`` is a ``GTensor(..., dtype=T.i32)`` over the fp8 bytes and + ``byte_off_i32`` is the byte offset of the first of the 8 bytes. + """ + dword_off = fx.Int32(arith.divui(_to_raw(byte_off_i32), _to_raw(fx.Int32(4)))) + v2 = i32_view.vec_load((dword_off,), vec_size=2) + return Vec(v2).bitcast(fx.Int64)[0].ir_value() + + +def fn_to_fnuz_i64(raw_i64): + """Map FN byte 0x80 (neg-zero) -> 0x00 in 8 packed fp8 bytes. + + An e4m3 FN (OCP, bias 7) operand encodes exactly 2x the value the gfx942 + FNUZ (bias 8) MFMA reads for the same byte, except FN -0 = 0x80 which is + FNUZ NaN; patch that byte to +0 and let the caller undo the 2x via the + kv_scale (ReLU positive-homogeneity keeps the result exact). + """ + lo_i32 = arith.TruncIOp(T.i32, raw_i64).result + hi_i64 = arith.ShRUIOp(raw_i64, arith.constant(32, type=T.i64)).result + hi_i32 = arith.TruncIOp(T.i32, hi_i64).result + + def _fix_i32(src): + result = arith.constant(0, type=T.i32) + for byte_idx in range_constexpr(4): + shift = arith.constant(byte_idx * 8, type=T.i32) + byte_val = arith.andi( + arith.shrui(src, shift), + arith.constant(0xFF, type=T.i32), + ) + is_0x80 = arith.cmpi( + arith.CmpIPredicate.eq, + byte_val, + arith.constant(0x80, type=T.i32), + ) + cleaned = arith.select( + is_0x80, + arith.constant(0, type=T.i32), + byte_val, + ) + result = arith.ori(result, arith.shli(cleaned, shift)) + return result + + lo_fix = _fix_i32(lo_i32) + hi_fix = _fix_i32(hi_i32) + lo_64 = arith.ExtUIOp(T.i64, lo_fix).result + hi_64 = arith.ShLIOp( + arith.ExtUIOp(T.i64, hi_fix).result, arith.constant(32, type=T.i64) + ).result + return arith.OrIOp(lo_64, hi_64).result + + +def make_out_row_view(logits, stride_i64, row_i32): + """1-D output GTensor for ``row_i32``; byte base computed in i64. + + The row's i64 byte offset (``row * stride * 4`` bytes for f32) goes into the + base pointer so the remaining column offset stays in i32. A 2-D (row, col) + view would compute ``row * stride + col`` in i32 and overflow past 2^31 + (~46k-square dense outputs / large paged ``max_model_len``), silently + mis-writing. + """ + _ri64 = arith.extui(T.i64, _to_raw(row_i32)) + _byte = arith.muli( + arith.muli(_ri64, stride_i64), arith.constant(4, type=T.i64) + ) + _idx = arith.index_cast(T.index, _byte) + return GTensor(logits, dtype=T.f32, shape=(-1,), static_bytes_offset_i64=_idx) + + +def mfma_head_reduce( + a_row, + b_col, + w_row, + kv_scale, + *, + m_tiles, + k_steps, + res_ty, + f32_0, + fm_fast, + mfma_fn=MFMA_FN, +): + """One column's logit: fp8 MFMA over heads, ReLU * weight sum, kv-scale, head reduce. + + ``a_row`` : ``a_row[mi][kk]`` packed-i64 Q fragments for this query row. + ``b_col`` : ``b_col[kk]`` packed-i64 K fragments for this KV column tile. + ``w_row`` : ``w_row[mi][ii]`` f32 head weights for this query row. + ``kv_scale``: f32 per-KV-token dequant scale (>= 0), hoisted out of the + head sum (ReLU is positive-homogeneous, so ReLU(s*x)=s*ReLU(x) + and the whole column sum is scaled once instead of per head). + + Returns the reduced f32 ``col_sum`` (raw ir value). The in-wave (width 64) + head reduce uses ``shuffle_xor`` offsets 16 and 32. + """ + col_sum = _to_raw(f32_0) + for mi in range_constexpr(m_tiles): + acc = Vec.filled(4, 0.0, fx.Float32) + for kk in range_constexpr(k_steps): + acc = mfma_fn(res_ty, [a_row[mi][kk], b_col[kk], acc, 0, 0, 0]) + for ii in range_constexpr(4): + score = Vec(acc)[ii].ir_value() + relu = arith.maximumf(score, _to_raw(f32_0)) + wsc = arith.MulFOp(relu, w_row[mi][ii], fastmath=fm_fast).result + col_sum = arith.AddFOp(col_sum, wsc, fastmath=fm_fast).result + col_sum = arith.MulFOp(col_sum, kv_scale, fastmath=fm_fast).result + + for sh in (16, 32): + peer = _to_raw(ArithValue(col_sum).shuffle_xor(sh, 64)) + col_sum = arith.AddFOp(col_sum, peer, fastmath=fm_fast).result + return col_sum diff --git a/aiter/ops/flydsl/kernels/fp8_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_mqa_logits.py index 2a61440cf84..ae2725c3373 100644 --- a/aiter/ops/flydsl/kernels/fp8_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_mqa_logits.py @@ -28,20 +28,25 @@ import flydsl.compiler as flyc import flydsl.expr as fx -from flydsl.expr import arith, range_constexpr, rocdl -from flydsl.expr.numeric import ArithValue +from flydsl.expr import arith, range_constexpr from flydsl.expr.typing import T from flydsl._mlir.dialects import scf from flydsl._mlir import ir from .tensor_shim import GTensor, _run_compiled, _to_raw - -Vec = fx.Vector - - -def _i32_add(a, b): - """i32 add (result stays Int32, not index type).""" - return fx.Int32(arith.addi(_to_raw(a), _to_raw(b))) +from ._mqa_logits_common import ( + MFMA_M, + MFMA_N, + MFMA_K, + DEFAULT_COMPILE_HINTS as _DEFAULT_COMPILE_HINTS, + Vec, + i32_add as _i32_add, + device_cu_count as _device_cu_count, + load_pack_i64, + fn_to_fnuz_i64, + make_out_row_view, + mfma_head_reduce, +) # Default KV tile width (columns processed per inner-loop iteration). @@ -50,20 +55,6 @@ def _i32_add(a, b): # below it the per-block Q/weight preload stops being amortized. _MIN_TILES_PER_SPLIT = 8 -_DEFAULT_COMPILE_HINTS = { - "waves_per_eu": 2, - "fast_fp_math": True, -} - - -@lru_cache(maxsize=8) -def _device_cu_count(device_index: int) -> int: - """Compute-unit count for a CUDA/HIP device (cached); 304 if unavailable.""" - try: - return torch.cuda.get_device_properties(device_index).multi_processor_count - except Exception: - return 304 - def _auto_num_splits( seq_len_padded: int, seq_len_kv: int, rpb: int, device_index: int @@ -121,12 +112,7 @@ def _build_kernel_mfma_r_w( WPB = waves_per_block MR_BLOCK_THREADS = 64 * WPB - # MFMA tile dims of the fp8 16x16x32 atom: MFMA_M x MFMA_N output tile, - # MFMA_K fp8 elements reduced per MFMA step. - MFMA_M = 16 - MFMA_N = 16 - MFMA_K = 32 - + # MFMA tile dims (MFMA_M/MFMA_N/MFMA_K) come from ._mqa_logits_common. assert H % MFMA_M == 0, f"num_heads={H} must be a multiple of MFMA_M={MFMA_M}" assert BKV % MFMA_N == 0, f"block_kv={BKV} must be a multiple of MFMA_N={MFMA_N}" assert D % MFMA_K == 0, f"head_size={D} must be a multiple of MFMA_K={MFMA_K}" @@ -141,7 +127,6 @@ def _build_kernel_mfma_r_w( N_TILES_PER_WAVE = N_TILES // WPB # column-tiles per wave fm_fast = arith.FastMathFlags.fast - mfma_fn = rocdl.mfma_f32_16x16x32_fp8_fp8 _cvt_tag = "" if convert_q_fn: @@ -199,59 +184,6 @@ def kernel( # (~46k-square dense outputs), silently mis-writing. _stride_i64 = arith.extui(T.i64, _to_raw(stride_logits_s)) - def _make_out_row_t(row_i32): - """1-D output GTensor for row_i32; byte base computed in i64.""" - _ri64 = arith.extui(T.i64, _to_raw(row_i32)) - _byte = arith.muli( - arith.muli(_ri64, _stride_i64), arith.constant(4, type=T.i64) - ) - _idx = arith.index_cast(T.index, _byte) - return GTensor( - logits, dtype=T.f32, shape=(-1,), static_bytes_offset_i64=_idx - ) - - def _load_pack_i64(i32_view, byte_off_i32): - dword_off = fx.Int32( - arith.divui(_to_raw(byte_off_i32), _to_raw(fx.Int32(4))) - ) - v2 = i32_view.vec_load((dword_off,), vec_size=2) - return Vec(v2).bitcast(fx.Int64)[0].ir_value() - - def _fn_to_fnuz_i64(raw_i64): - """Map FN byte 0x80 (neg-zero) -> 0x00 in 8 packed fp8 bytes.""" - lo_i32 = arith.TruncIOp(T.i32, raw_i64).result - hi_i64 = arith.ShRUIOp(raw_i64, arith.constant(32, type=T.i64)).result - hi_i32 = arith.TruncIOp(T.i32, hi_i64).result - - def _fix_i32(src): - result = arith.constant(0, type=T.i32) - for byte_idx in range_constexpr(4): - shift = arith.constant(byte_idx * 8, type=T.i32) - byte_val = arith.andi( - arith.shrui(src, shift), - arith.constant(0xFF, type=T.i32), - ) - is_0x80 = arith.cmpi( - arith.CmpIPredicate.eq, - byte_val, - arith.constant(0x80, type=T.i32), - ) - cleaned = arith.select( - is_0x80, - arith.constant(0, type=T.i32), - byte_val, - ) - result = arith.ori(result, arith.shli(cleaned, shift)) - return result - - lo_fix = _fix_i32(lo_i32) - hi_fix = _fix_i32(hi_i32) - lo_64 = arith.ExtUIOp(T.i64, lo_fix).result - hi_64 = arith.ShLIOp( - arith.ExtUIOp(T.i64, hi_fix).result, arith.constant(32, type=T.i64) - ).result - return arith.OrIOp(lo_64, hi_64).result - # ---- Preload window bounds, Q frags, and weights for all RPB rows ---- # A-operand layout is per in-wave lane, so `lane` (not `tid`) indexes Q. starts = [None] * RPB @@ -277,8 +209,8 @@ def _fix_i32(src): base_a = fx.Int32(arith.muli(_to_raw(row_h), _to_raw(fx.Int32(D)))) for kk in range_constexpr(K_STEPS): d_a = _i32_add(fx.Int32(kk * MFMA_K), lane8) - raw = _load_pack_i64(q_i32, _i32_add(base_a, d_a)) - row_a[mi][kk] = _fn_to_fnuz_i64(raw) if convert_q_fn else raw + raw = load_pack_i64(q_i32, _i32_add(base_a, d_a)) + row_a[mi][kk] = fn_to_fnuz_i64(raw) if convert_q_fn else raw a_packs[j] = row_a # weights[row, h] per (mi, ii): head = mi*MFMA_M + lane_div_N*4 + ii @@ -363,43 +295,29 @@ def _fix_i32(src): ) for kk in range_constexpr(K_STEPS): d_b = _i32_add(fx.Int32(kk * MFMA_K), lane8) - raw = _load_pack_i64(kv_i32, _i32_add(base_b, d_b)) - b_packs[ni][kk] = _fn_to_fnuz_i64(raw) if convert_kv_fn else raw + raw = load_pack_i64(kv_i32, _i32_add(base_b, d_b)) + b_packs[ni][kk] = fn_to_fnuz_i64(raw) if convert_kv_fn else raw # ---- Per-row MFMA + epilogue (inner loop over RPB rows) ---- for j in range_constexpr(RPB): row = _i32_add(r0, fx.Int32(j)) - out_row_t = _make_out_row_t(row) + out_row_t = make_out_row_view(logits, _stride_i64, row) for ni in range_constexpr(N_TILES_PER_WAVE): col = cols[ni] kv_scale = kv_scales_tile[ni] - col_sum = _to_raw(f32_0) - for mi in range_constexpr(M_TILES): - acc = Vec.filled(4, 0.0, fx.Float32) - for kk in range_constexpr(K_STEPS): - acc = mfma_fn( - mfma_res_ty, - [a_packs[j][mi][kk], b_packs[ni][kk], acc, 0, 0, 0], - ) - # kv_scale (>=0) is hoisted out of the head sum: ReLU is - # positive-homogeneous, so ReLU(s*x)=s*ReLU(x) and the - # whole column sum is scaled once (below) instead of every - # head term -- drops M_TILES*4 muls to one. - for ii in range_constexpr(4): - score = Vec(acc)[ii].ir_value() - relu = arith.maximumf(score, _to_raw(f32_0)) - wsc = arith.MulFOp( - relu, w_frag[j][mi][ii], fastmath=fm_fast - ).result - col_sum = arith.AddFOp( - col_sum, wsc, fastmath=fm_fast - ).result - col_sum = arith.MulFOp(col_sum, kv_scale, fastmath=fm_fast).result - - # Head-reduce within the wave (width=64): shuffle_xor offsets 16, 32. - for sh in [16, 32]: - peer = _to_raw(ArithValue(col_sum).shuffle_xor(sh, 64)) - col_sum = arith.AddFOp(col_sum, peer, fastmath=fm_fast).result + # fp8 MFMA over heads + ReLU*weight head-sum + kv-scale hoist + # + in-wave shuffle_xor head reduce (see ._mqa_logits_common). + col_sum = mfma_head_reduce( + a_packs[j], + b_packs[ni], + w_frag[j], + kv_scale, + m_tiles=M_TILES, + k_steps=K_STEPS, + res_ty=mfma_res_ty, + f32_0=f32_0, + fm_fast=fm_fast, + ) # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. # `col >= start` is required: the tile loop is BKV-aligned @@ -518,7 +436,6 @@ def compile_fp8_mqa_logits( num_heads: int, head_size: int, block_kv: int = _BLOCK_KV, - paged: bool = False, variant: str = DEFAULT_VARIANT, convert_q_fn: bool = False, convert_kv_fn: bool = False, @@ -528,13 +445,11 @@ def compile_fp8_mqa_logits( ``num_heads``/``head_size`` are compile-time constants (powers of two, D in {64, 128}); ``variant`` is an ``mfma_r_w`` tag (see ``KERNEL_VARIANTS``); ``convert_q_fn``/``convert_kv_fn`` mark an FP8 FN - operand whose -0 (0x80) byte the kernel patches to FNUZ +0. ``paged`` is - reserved for a future variant and must be False. + operand whose -0 (0x80) byte the kernel patches to FNUZ +0. + + The paged KV variant lives in its own module + (``fp8_paged_mqa_logits.flydsl_fp8_paged_mqa_logits``). """ - if paged: - raise NotImplementedError( - "Paged FlyDSL fp8_mqa_logits is Phase 2 and not implemented yet." - ) if variant not in _VARIANT_BUILDERS: raise ValueError( f"unknown fp8_mqa_logits variant {variant!r}; " @@ -622,7 +537,6 @@ def flydsl_fp8_mqa_logits( num_heads=num_heads, head_size=head_size, block_kv=_BLOCK_KV, - paged=False, variant=variant, convert_q_fn=convert_q_fn, convert_kv_fn=convert_kv_fn, diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py new file mode 100644 index 00000000000..92e93e6d03f --- /dev/null +++ b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py @@ -0,0 +1,596 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Paged FP8 MQA logits (DeepSeek lightning indexer, decode) -- FlyDSL gfx942/gfx950. + +Compute for each decode batch element ``b`` and next-n query slot ``n`` (query +row ``b*next_n + n``) and KV logical position ``p``:: + + logits[b*next_n+n, p] = sum_h ReLU() * weights[b*next_n+n, h] * kv_scale_phys(p) + +masked by the causal rule ``p <= context_lens[b] - next_n + n`` (positions past +the query token and ``p >= context_lens[b]`` stay ``-inf``). ``K_phys(p)`` / +``kv_scale_phys(p)`` are gathered from a **paged** cache: logical position ``p`` +maps through the block table (``kv_indices``) to a physical block, and each +token's fp8 K bytes and its f32 dequant scale are **co-packed** in that block. + +Supported scope (this file): ``KVBlockSize == 1`` with SplitKV KV-column +parallelism. The public ``flydsl_fp8_paged_mqa_logits`` mirrors the tensor +contract of the Triton +``aiter.ops.triton.attention.pa_mqa_logits.deepgemm_fp8_paged_mqa_logits`` so the +two are interchangeable in tests/benchmarks. ``KVBlockSize > 1`` and preshuffle +are not supported and raise ``NotImplementedError``. + +The fp8 16x16x32 MFMA compute, ReLU*weight head-sum + kv-scale hoist, in-wave +``shuffle_xor`` head reduce, the fp8 dword-pack load, the FN->FNUZ byte patch, and +the i64 byte-base per-row output view are all shared with the dense kernel via +``._mqa_logits_common`` (no duplication). +""" + +# No `from __future__ import annotations`: FlyDSL arg typing needs real +# annotation objects, not PEP 563 strings. + +import os +from functools import lru_cache + +import torch + +from aiter.jit.utils.chip_info import get_gfx + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith, range_constexpr +from flydsl.expr.typing import T +from flydsl._mlir.dialects import scf +from flydsl._mlir import ir + +from .tensor_shim import GTensor, _run_compiled, _to_raw +from ._mqa_logits_common import ( + MFMA_M, + MFMA_N, + MFMA_K, + DEFAULT_COMPILE_HINTS, + Vec, + i32_add as _i32_add, + device_cu_count, + load_pack_i64, + fn_to_fnuz_i64, + make_out_row_view, + mfma_head_reduce, +) + +def _auto_split_kv(batch_size, next_n, wave_per_eu, device_index, total_cu=None): + """Production host SplitKV formula (mirrors deepgemm_fp8_paged_mqa_logits): + + SplitKV = ((max(1, TotalCuCount // (batch*next_n)) + 4) // 5 * 5) * WavePerEU + + Fills the device for small decode grids where the batch*next_n row grid alone + leaves most CUs idle. Returns >= 1. + """ + tile_q_count = max(1, batch_size * next_n) + if total_cu is None: + total_cu = device_cu_count(device_index) + split_kv = ((max(1, total_cu // tile_q_count) + 4) // 5 * 5) * wave_per_eu + return max(1, int(split_kv)) + +# Default KV tile width (columns processed per MFMA inner-loop iteration). +_BLOCK_KV = 128 + + +def _build_paged_kernel( + *, + num_heads: int, + head_size: int, + block_kv: int, + waves_per_block: int, + scale_mul: float = 1.0, + convert_q_fn: bool = False, + convert_kv_fn: bool = False, +): + """Paged MQA-logits kernel (KVBlockSize==1) with SplitKV column parallelism. + + One thread block owns one ``(batch, next_n)`` query row and walks the whole + ``[0, context_length)`` KV window in ``BKV``-wide tiles. ``waves_per_block`` + waves each own a disjoint slice of the ``BKV/MFMA_N`` column tiles (no + cross-wave sharing / barrier), mirroring the dense kernel's wave split. + + ``scale_mul`` folds the FN->FNUZ 2x-per-converted-operand compensation into + the (co-packed) per-token scale at compile time (ReLU positive-homogeneity). + """ + H = num_heads + D = head_size + BKV = block_kv + WPB = waves_per_block + MR_BLOCK_THREADS = 64 * WPB + + assert H % MFMA_M == 0, f"num_heads={H} must be a multiple of MFMA_M={MFMA_M}" + assert BKV % MFMA_N == 0, f"block_kv={BKV} must be a multiple of MFMA_N={MFMA_N}" + assert D % MFMA_K == 0, f"head_size={D} must be a multiple of MFMA_K={MFMA_K}" + assert WPB >= 1, "waves_per_block must be >= 1" + N_TILES = BKV // MFMA_N + assert ( + N_TILES % WPB == 0 + ), f"BKV/MFMA_N={N_TILES} must be divisible by waves_per_block={WPB}" + M_TILES = H // MFMA_M + K_STEPS = D // MFMA_K + N_TILES_PER_WAVE = N_TILES // WPB + + fm_fast = arith.FastMathFlags.fast + D_over_4 = D // 4 # f32-element offset of the co-packed scale within a token + + _cvt_tag = "" + if convert_q_fn: + _cvt_tag += "_cq" + if convert_kv_fn: + _cvt_tag += "_ck" + _kname = f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_w{WPB}{_cvt_tag}_flydsl" + + @flyc.kernel(name=_kname, known_block_size=[MR_BLOCK_THREADS, 1, 1]) + def kernel( + Q: fx.Tensor, # [batch, next_n, H, D] fp8 (bytes passed raw) + KV_cache: fx.Tensor, # [num_blocks, 1, 1, index_dim] uint8 co-packed + weights: fx.Tensor, # [batch*next_n, H] f32 + out_logits: fx.Tensor, # [batch*next_n, max_model_len] f32 (-inf prefilled) + context_lens: fx.Tensor, # [batch] i32 + kv_indices: fx.Tensor, # [batch, max_block_len] i32 (block table) + next_n: fx.Int32, + batch_size: fx.Int32, + split_kv: fx.Int32, # KV-column splits (grid = split_kv*batch*next_n) + stride_q_batch: fx.Int32, # fp8 elems (== bytes) + stride_q_next_n: fx.Int32, + stride_q_heads: fx.Int32, + index_dim: fx.Int32, # cache row width in bytes (D + 4 + pad) + max_block_len: fx.Int32, # kv_indices row width + stride_out: fx.Int32, # out_logits.stride(0) == max_model_len + ): + f32_0 = arith.constant(0.0, type=T.f32) + mfma_res_ty = Vec.make_type(4, fx.Float32) + scale_mul_c = arith.constant(float(scale_mul), type=T.f32) + + tid = fx.thread_idx.x + bid = fx.block_idx.x + + # ChunkQ==H => grid = split_kv * batch * next_n (split outermost), + # mirroring the Triton _deepgemm_fp8_paged_mqa_logits pid decomposition. + pid_next_n = fx.Int32(arith.remui(_to_raw(bid), _to_raw(next_n))) + _rem = fx.Int32(arith.divui(_to_raw(bid), _to_raw(next_n))) + pid_batch = fx.Int32(arith.remui(_to_raw(_rem), _to_raw(batch_size))) + pid_split_kv = fx.Int32(arith.divui(_to_raw(_rem), _to_raw(batch_size))) + + wave = fx.Int32(arith.divui(_to_raw(tid), _to_raw(fx.Int32(64)))) + lane = fx.Int32(arith.remui(_to_raw(tid), _to_raw(fx.Int32(64)))) + lane_div_N = fx.Int32(arith.divui(_to_raw(lane), _to_raw(fx.Int32(MFMA_N)))) + lane_mod_N = fx.Int32(arith.remui(_to_raw(lane), _to_raw(fx.Int32(MFMA_N)))) + lane8 = fx.Int32(arith.muli(_to_raw(lane_div_N), _to_raw(fx.Int32(8)))) + + # fp8 operands read 8 bytes at a time as 2 i32 dwords (v8i8 buffer_load + # fails to lower on gfx942), then bitcast to i64 for the MFMA. + q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) + # Uniform-base views over the co-packed cache: per-token gather is done + # via the (per-lane) byte offset, NOT the buffer base -- a per-lane base + # can't ride a scalar buffer descriptor. The byte offset stays i32 + # (assumes num_blocks*index_dim < 2^31; see host assert); a shared cache + # pool large enough to exceed that would need an i64/global-load gather. + kv_i32 = GTensor(KV_cache, dtype=T.i32, shape=(-1,)) + kv_f32 = GTensor(KV_cache, dtype=T.f32, shape=(-1,)) + w_t = GTensor(weights, dtype=T.f32, shape=(-1, H)) + cl_t = GTensor(context_lens, dtype=T.i32, shape=(-1,)) + ind_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) + + context_length = fx.Int32(cl_t[pid_batch]) + # Inclusive causal upper bound: p <= context_length - next_n + pid_next_n. + q_limit = fx.Int32( + arith.addi( + arith.subi(_to_raw(context_length), _to_raw(next_n)), + _to_raw(pid_next_n), + ) + ) + + out_row = _i32_add( + fx.Int32(arith.muli(_to_raw(pid_batch), _to_raw(next_n))), pid_next_n + ) + stride_out_i64 = arith.extui(T.i64, _to_raw(stride_out)) + out_row_t = make_out_row_view(out_logits, stride_out_i64, out_row) + + # ---- Preload Q frags + weights for the single query row ---- + # A-operand layout is per in-wave lane, so `lane` (not `tid`) indexes Q. + q_row_base = _i32_add( + fx.Int32(arith.muli(_to_raw(pid_batch), _to_raw(stride_q_batch))), + fx.Int32(arith.muli(_to_raw(pid_next_n), _to_raw(stride_q_next_n))), + ) + a_pack = [[None] * K_STEPS for _ in range_constexpr(M_TILES)] + for mi in range_constexpr(M_TILES): + h_a = _i32_add(fx.Int32(mi * MFMA_M), lane_mod_N) + base_a = _i32_add( + q_row_base, + fx.Int32(arith.muli(_to_raw(h_a), _to_raw(stride_q_heads))), + ) + for kk in range_constexpr(K_STEPS): + d_a = _i32_add(fx.Int32(kk * MFMA_K), lane8) + raw = load_pack_i64(q_i32, _i32_add(base_a, d_a)) + a_pack[mi][kk] = fn_to_fnuz_i64(raw) if convert_q_fn else raw + + # weights[out_row, h] per (mi, ii): head = mi*MFMA_M + lane_div_N*4 + ii + w_frag = [[None] * 4 for _ in range_constexpr(M_TILES)] + for mi in range_constexpr(M_TILES): + for ii in range_constexpr(4): + h_w = _i32_add( + fx.Int32(mi * MFMA_M), + _i32_add( + fx.Int32( + arith.muli(_to_raw(lane_div_N), _to_raw(fx.Int32(4))) + ), + fx.Int32(ii), + ), + ) + w_frag[mi][ii] = _to_raw(fx.Float32(w_t[out_row, h_w])) + + # ---- SplitKV: this CTA owns a disjoint, BKV-aligned slice of the KV + # window. total_tiles = ceil(ctx/BKV); each split owns + # ceil(total_tiles/split_kv) contiguous tiles => slices are gap-free and + # disjoint, so every logits column has exactly one writer (pure + # parallelism across n, no cross-CTA reduction). split_kv==1 collapses to + # the full [0, ceil(ctx/BKV)*BKV) window. Splits past ctx early-exit + # (tile_lo_col >= tile_hi_col => the loop runs 0 iterations). ---- + context_chunk_num = arith.ceildivui( + _to_raw(context_length), _to_raw(fx.Int32(BKV)) + ) + split_chunk_num = arith.ceildivui(context_chunk_num, _to_raw(split_kv)) + full_end = arith.muli(context_chunk_num, _to_raw(fx.Int32(BKV))) + split_cols = arith.muli(split_chunk_num, _to_raw(fx.Int32(BKV))) + tile_lo_col = arith.muli(_to_raw(pid_split_kv), split_cols) + tile_hi_col = arith.minsi(arith.addi(tile_lo_col, split_cols), full_end) + ctx_m1 = arith.subi(_to_raw(context_length), _to_raw(fx.Int32(1))) + + tile_lo = _to_raw(fx.Index(fx.Int32(tile_lo_col))) + tile_hi = _to_raw(fx.Index(fx.Int32(tile_hi_col))) + tile_step = _to_raw(fx.Index(fx.Int32(BKV))) + tile_loop = scf.ForOp(tile_lo, tile_hi, tile_step, []) + with ir.InsertionPoint(tile_loop.body): + col0 = fx.Int32(arith.index_cast(T.i32, tile_loop.induction_variable)) + wave_ni_base = fx.Int32( + arith.muli(_to_raw(wave), _to_raw(fx.Int32(N_TILES_PER_WAVE))) + ) + for ni in range_constexpr(N_TILES_PER_WAVE): + abs_ni = _i32_add(wave_ni_base, fx.Int32(ni)) + col = _i32_add( + _i32_add( + col0, + fx.Int32( + arith.muli(_to_raw(abs_ni), _to_raw(fx.Int32(MFMA_N))) + ), + ), + lane_mod_N, + ) + # Clamp the gather index to a valid token (mask handles the rest). + col_c = fx.Int32(arith.minsi(_to_raw(col), ctx_m1)) + ind_off = _i32_add( + fx.Int32(arith.muli(_to_raw(pid_batch), _to_raw(max_block_len))), + col_c, + ) + physical = fx.Int32(ind_t[ind_off]) + tok_byte = fx.Int32( + arith.muli(_to_raw(physical), _to_raw(index_dim)) + ) + + b_col = [None] * K_STEPS + for kk in range_constexpr(K_STEPS): + d_b = _i32_add(fx.Int32(kk * MFMA_K), lane8) + raw = load_pack_i64(kv_i32, _i32_add(tok_byte, d_b)) + b_col[kk] = fn_to_fnuz_i64(raw) if convert_kv_fn else raw + + # Co-packed per-token scale: f32 at byte tok_byte + D. + scale_dword = fx.Int32( + arith.divui( + _to_raw(_i32_add(tok_byte, fx.Int32(D))), + _to_raw(fx.Int32(4)), + ) + ) + kv_scale = _to_raw(fx.Float32(kv_f32[scale_dword])) + if scale_mul != 1.0: + kv_scale = arith.MulFOp( + kv_scale, scale_mul_c, fastmath=fm_fast + ).result + + col_sum = mfma_head_reduce( + a_pack, + b_col, + w_frag, + kv_scale, + m_tiles=M_TILES, + k_steps=K_STEPS, + res_ty=mfma_res_ty, + f32_0=f32_0, + fm_fast=fm_fast, + ) + + # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. + # Write only in-causal columns; masked positions keep the + # caller's -inf prefill (matches clean_logits semantics). + is_writer = arith.andi( + _to_raw( + arith.cmpi( + arith.CmpIPredicate.eq, + _to_raw(lane_div_N), + _to_raw(fx.Int32(0)), + ) + ), + _to_raw( + arith.cmpi( + arith.CmpIPredicate.sle, + _to_raw(col), + _to_raw(q_limit), + ) + ), + ) + with ir.InsertionPoint(scf.IfOp(is_writer).then_block): + out_row_t[col] = fx.Float32(col_sum) + scf.YieldOp([]) + + scf.YieldOp([]) + + @flyc.jit + def launch_fp8_paged_mqa_logits( + Q: fx.Tensor, + KV_cache: fx.Tensor, + weights: fx.Tensor, + out_logits: fx.Tensor, + context_lens: fx.Tensor, + kv_indices: fx.Tensor, + grid_blocks: fx.Int32, + next_n: fx.Int32, + batch_size: fx.Int32, + split_kv: fx.Int32, + stride_q_batch: fx.Int32, + stride_q_next_n: fx.Int32, + stride_q_heads: fx.Int32, + index_dim: fx.Int32, + max_block_len: fx.Int32, + stride_out: fx.Int32, + stream: fx.Stream = fx.Stream(None), + ): + gx = arith.index_cast(T.index, _to_raw(grid_blocks)) + kernel._func.__name__ = _kname + kernel( + Q, + KV_cache, + weights, + out_logits, + context_lens, + kv_indices, + next_n, + batch_size, + split_kv, + stride_q_batch, + stride_q_next_n, + stride_q_heads, + index_dim, + max_block_len, + stride_out, + ).launch( + grid=(gx, 1, 1), block=(MR_BLOCK_THREADS, 1, 1), stream=stream + ) + + return launch_fp8_paged_mqa_logits + + +# Kernel variants: single-token-per-block, wave-split only ("paged_w"). +# WPB must divide the column-tile count BKV/16 (=8 at the default BKV=128). +def _mk_builder(wpb): + return lambda **kw: _build_paged_kernel(**kw, waves_per_block=wpb) + + +_VARIANT_BUILDERS = {f"paged_w{w}": _mk_builder(w) for w in (1, 2, 4)} +KERNEL_VARIANTS = tuple(_VARIANT_BUILDERS.keys()) +DEFAULT_VARIANT = "paged_w4" + + +def _auto_variant(batch_size, next_n, heads): + """Default variant selector: constant WPB=4 (8 column tiles / 4 waves = 2 + tiles/wave). Shape-adaptive variant selection is not implemented.""" + return DEFAULT_VARIANT + + +def _resolve_variant(variant, batch_size, next_n, heads): + tag = ( + variant + or os.environ.get("FLYDSL_FP8_PAGED_MQA_LOGITS_VARIANT") + or _auto_variant(batch_size, next_n, heads) + ) + if tag not in _VARIANT_BUILDERS: + raise ValueError( + f"unknown fp8_paged_mqa_logits variant {tag!r}; " + f"available: {list(KERNEL_VARIANTS)}" + ) + return tag + + +@lru_cache(maxsize=32) +def compile_fp8_paged_mqa_logits( + *, + num_heads: int, + head_size: int, + block_kv: int = _BLOCK_KV, + variant: str = DEFAULT_VARIANT, + scale_mul: float = 1.0, + convert_q_fn: bool = False, + convert_kv_fn: bool = False, +): + """Return a cached, compiled FlyDSL paged launcher for the given config. + + ``num_heads``/``head_size`` are compile-time constants; ``variant`` is a + ``paged_w`` tag; ``convert_q_fn``/``convert_kv_fn`` mark an FP8 FN + operand whose -0 (0x80) byte the kernel patches to FNUZ +0, and ``scale_mul`` + folds the 2x-per-converted-operand compensation into the co-packed scale. + """ + if variant not in _VARIANT_BUILDERS: + raise ValueError( + f"unknown fp8_paged_mqa_logits variant {variant!r}; " + f"available: {list(KERNEL_VARIANTS)}" + ) + launcher = _VARIANT_BUILDERS[variant]( + num_heads=num_heads, + head_size=head_size, + block_kv=block_kv, + scale_mul=scale_mul, + convert_q_fn=convert_q_fn, + convert_kv_fn=convert_kv_fn, + ) + launcher.compile_hints = dict(DEFAULT_COMPILE_HINTS) + return launcher + + +def flydsl_fp8_paged_mqa_logits( + q_fp8, + kv_cache, + weights, + out_logits, + context_lens, + kv_indices, + max_model_len, + *, + Preshuffle=False, + KVBlockSize=1, + ChunkK=256, + SplitKV=None, + WavePerEU=2, + TotalCuCount=None, + variant=None, + stream=None, +): + """FlyDSL paged FP8 MQA logits (decode) -- KVBlockSize==1 with SplitKV. + + Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. + + q_fp8: [batch, next_n, heads, hidden_dim], dtype float8 (e4m3 fn/fnuz) + kv_cache: [num_blocks, KVBlockSize, 1, index_dim], dtype uint8, co-packed + fp8 K bytes (first KVBlockSize*hidden_dim) then f32 dequant + scales; index_dim == hidden_dim + 4 (+ optional 16B padding). + weights: [batch*next_n, heads], dtype float32 + out_logits: [batch*next_n, max_model_len], dtype float32. MUST be prefilled + with -inf by the caller; the kernel writes only in-window + (causal) positions and leaves masked positions untouched. + context_lens: [batch], dtype int32 + kv_indices: [batch, max_block_len], dtype int32 block table (physical block + per logical position; KVBlockSize==1 => per-token page table). + max_model_len: int, out_logits column count. + SplitKV: KV-column split count (grid = split_kv*batch*next_n). None => + the production host formula (fills the device on small decode + grids); pass an int to override, 1 disables splitting. Splits + own disjoint, gap-free BKV-aligned column ranges (one writer + per column, no cross-CTA reduction). + WavePerEU/TotalCuCount: inputs to the auto-SplitKV formula. + ChunkK: accepted for Triton-contract parity; ignored here (the FlyDSL + kernel tiles the KV window at a fixed ``block_kv`` width). + + Returns the same ``out_logits`` tensor (written in place). + """ + if Preshuffle: + raise NotImplementedError( + "Preshuffle is not supported by the FlyDSL paged fp8_mqa_logits kernel." + ) + if KVBlockSize != 1: + raise NotImplementedError( + "FlyDSL paged fp8_mqa_logits supports KVBlockSize==1 only " + f"(got KVBlockSize={KVBlockSize})." + ) + + batch_size, next_n, num_heads, head_size = q_fp8.shape + assert num_heads & (num_heads - 1) == 0, "num q. heads should be power of 2." + assert head_size & (head_size - 1) == 0, "head size should be power of 2." + + num_blocks, block_size, one, index_dim = kv_cache.shape + assert block_size == 1, ( + f"kv_cache KVBlockSize dim must be 1; got {block_size}." + ) + assert one == 1, f"kv_cache head dim must be 1; got {one}." + assert index_dim >= head_size + 4, ( + f"index_dim={index_dim} must hold {head_size} fp8 bytes + 4 scale bytes." + ) + assert kv_cache.dtype == torch.uint8, ( + f"kv_cache must be uint8 co-packed bytes; got {kv_cache.dtype}." + ) + + # i32 gather-offset ceiling: the per-token byte base physical*index_dim is + # computed in i32 (it rides the buffer voffset), so the whole cache pool must + # be addressable in i32 bytes. A larger pool would need an i64/global gather. + assert num_blocks * index_dim < 2**31, ( + f"num_blocks*index_dim={num_blocks * index_dim} exceeds the i32 " + f"gather-offset limit (2^31); the paged kernel needs an i64 gather path " + f"for a cache pool this large." + ) + + _, max_block_len = kv_indices.shape + + # FlyDSL's DLPack adaptor rejects 0-dim tensors; keep logical ranks. + context_lens = context_lens.reshape(batch_size) + weights = weights.reshape(batch_size * next_n, num_heads) + kv_indices = kv_indices.reshape(batch_size, max_block_len) + + _fnuz = torch.float8_e4m3fnuz + _fn = torch.float8_e4m3fn + assert q_fp8.dtype in (_fnuz, _fn), ( + f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" + ) + + # gfx942 fp8 MFMA reads operands as e4m3 FNUZ (bias 8). Q may arrive as FN + # (OCP, bias 7); patch its 0x80 byte and undo the implied 2x via scale_mul. + # The co-packed KV in the cache is arch-native fp8 (fnuz on gfx942, fn on + # gfx950), so it is never converted. + gfx = get_gfx() + convert_q_fn = gfx == "gfx942" and q_fp8.dtype == _fn + convert_kv_fn = False + scale_mul = 2.0 if convert_q_fn else 1.0 + + variant = _resolve_variant(variant, batch_size, next_n, num_heads) + + launcher = compile_fp8_paged_mqa_logits( + num_heads=num_heads, + head_size=head_size, + block_kv=_BLOCK_KV, + variant=variant, + scale_mul=scale_mul, + convert_q_fn=convert_q_fn, + convert_kv_fn=convert_kv_fn, + ) + + # Co-packed cache -> raw uint8 byte view [num_blocks * index_dim]. + kv_bytes = kv_cache.reshape(-1) + + # KV-column splits (one lever): fill the device when the batch*next_n row + # grid is small. Auto == production host formula; overridable. + if SplitKV is None: + split_kv = _auto_split_kv( + batch_size, next_n, WavePerEU, q_fp8.device.index, TotalCuCount + ) + else: + split_kv = max(1, int(SplitKV)) + + grid_blocks = batch_size * next_n * split_kv + + if stream is None: + stream = torch.cuda.current_stream() + + with torch.cuda.device(q_fp8.device.index): + _run_compiled( + launcher, + q_fp8, + kv_bytes, + weights, + out_logits, + context_lens, + kv_indices, + int(grid_blocks), + int(next_n), + int(batch_size), + int(split_kv), + int(q_fp8.stride(0)), + int(q_fp8.stride(1)), + int(q_fp8.stride(2)), + int(index_dim), + int(max_block_len), + int(out_logits.stride(0)), + stream, + ) + + return out_logits diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py new file mode 100644 index 00000000000..0a25a4d6da7 --- /dev/null +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -0,0 +1,304 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness + perf sweep for the FlyDSL paged FP8 MQA-logits (decode) kernel. + +Mirrors the aiter-op-test standard used by ``test_flydsl_fp8_mqa_logits.py`` and +the paged benchmark ``op_tests/op_benchmarks/triton/bench_deepgemm_attention.py``. +The correctness reference is a torch port of vLLM's ``fp8_paged_mqa_logits_torch`` +(dequantizes the co-packed fp8 cache and applies the causal mask); it is copied +here so the test does not depend on vLLM. Correctness gate: exact ``-inf``-mask +match + ``calc_diff < 1e-3`` (tolerances are NOT widened). +""" + +import argparse +import itertools +import random + +import aiter +import pandas as pd +import torch +from aiter import dtypes +from aiter.test_common import benchmark, checkAllclose, run_perftest +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype + +torch.set_default_device("cuda") + +# gfx942 is the primary target; gfx950 (native fp8) is also runnable. +SUPPORTED_GFX = ["gfx942", "gfx950"] +# "fnuz" == the arch-native e4m3 (fnuz on gfx942, fn on gfx950); "fn" forces the +# OCP e4m3fn operand (exercises the gfx942 FN->FNUZ patch). Mirrors the dense +# test's dtype aliasing. +_E4M3_NATIVE = get_fp8_e4m3_dtype() +DTYPE_MAP = {"fnuz": _E4M3_NATIVE, "fn": torch.float8_e4m3fn} + +try: + from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits +except ImportError: + flydsl_fp8_paged_mqa_logits = None + + +def calc_diff(x, y): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + return 1 - 2 * (x * y).sum() / denominator + + +def kv_cache_cast_to_fp8(x, fp8_dtype): + """Co-pack a bf16 KV cache into the fp8+scale byte layout (KVBlockSize>=1). + + Layout per block-row: KVBlockSize*head_dim fp8 bytes, then KVBlockSize f32 + (4-byte) per-token scales. No 16B padding (Phase 1). Mirrors the paged + benchmark's builder. + """ + num_blocks, block_size, num_heads, head_dim = x.shape + assert num_heads == 1 + x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) + sf = x_amax / 240.0 + x_scaled = (x * (1.0 / sf)).to(fp8_dtype) + x_fp8 = torch.empty( + (num_blocks, block_size * (head_dim + 4)), + device=x.device, + dtype=torch.uint8, + ) + x_fp8[:, : block_size * head_dim] = x_scaled.view( + num_blocks, block_size * head_dim + ).view(dtype=torch.uint8) + x_fp8[:, block_size * head_dim : block_size * head_dim + 4 * block_size] = sf.view( + num_blocks, block_size + ).view(dtype=torch.uint8) + return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) + + +def ref_fp8_paged_mqa_logits( + q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype +): + """Torch reference (vectorized port of vLLM ``fp8_paged_mqa_logits_torch``). + + Dequantizes the co-packed fp8 cache (``fp8_bytes.float() * scale``) so it + matches what the kernel computes (kv-scale folded into K), then: + logits[b*next_n+n, p] = sum_h ReLU() * weights[.., h] + with the causal mask ``p <= context_len - next_n + n``. The inner per-token + loop of the original is replaced by a single gather + einsum per batch + element (block_size == 1), which is orders of magnitude faster at long + context while numerically identical. + """ + batch_size, next_n, heads, dim = q.size() + kvv, scale = kv_cache_fp8[..., :dim], kv_cache_fp8[..., dim:] + scale = scale.contiguous().view(torch.float) + qf = q.float() + # dequantized K per physical block: [num_blocks, dim] (block_size == 1) + kvf = (kvv.view(fp8_dtype).float() * scale).view(kv_cache_fp8.shape[0], dim) + logits = torch.full( + [batch_size * next_n, max_model_len], + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + for i in range(batch_size): + context_len = int(context_lens[i].item()) + if context_len == 0: + continue + pages = block_tables[i, :context_len] # [ctx] physical token ids + kx = kvf[pages] # [ctx, dim] dequantized K + s = torch.einsum("nhd,pd->nhp", qf[i], kx) # [next_n, heads, ctx] + s = torch.relu(s) + wl = weights[i * next_n : (i + 1) * next_n, :] # [next_n, heads] + s = (s * wl[:, :, None]).sum(dim=1) # [next_n, ctx] + p = torch.arange(context_len, device=q.device) + q_lim = ( + context_len - next_n + torch.arange(next_n, device=q.device) + ).unsqueeze(1) # [next_n, 1] + s = torch.where(p[None, :] <= q_lim, s, float("-inf")) + logits[i * next_n : (i + 1) * next_n, :context_len] = s + return logits + + +def _build_inputs(batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, seed=0): + torch.manual_seed(seed) + random.seed(seed) + fp8_dtype = get_fp8_e4m3_dtype() + + max_model_len = 2 * avg_kv_length + num_blocks = max_model_len # KVBlockSize == 1 + + lo = max(1, int((1 - 0.5) * avg_kv_length)) + hi = int((1 + 0.5) * avg_kv_length) + 1 + context_lens = torch.randint(lo, hi, (batch_size,)).cuda().to(torch.int32) + # decode with MTP needs at least next_n tokens of context. + context_lens = torch.clamp(context_lens, min=next_n) + + q = torch.randn((batch_size, next_n, heads, head_dim), dtype=torch.bfloat16) + kv_cache = torch.randn((num_blocks, 1, 1, head_dim), dtype=torch.bfloat16) + weights = torch.randn((batch_size * next_n, heads), dtype=torch.float32) + + max_block_len = int(context_lens.max().item()) # blocksize == 1 + block_tables = torch.zeros( + (batch_size, max_block_len), device="cuda", dtype=torch.int32 + ) + pool = list(range(num_blocks)) + random.shuffle(pool) + counter = 0 + for i in range(batch_size): + for j in range(int(context_lens[i].item())): + block_tables[i][j] = pool[counter % num_blocks] + counter += 1 + + q_fp8 = q.to(q_dtype) + kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, fp8_dtype) + return ( + q, + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + ) + + +@benchmark() +def test_fp8_paged_mqa_logits( + batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, split_kv=0 +): + # split_kv == 0 -> auto (production host formula); else an explicit override + # (1 disables splitting). Both must be correctness-identical. + _split_kv = None if split_kv == 0 else split_kv + ( + q, + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + ) = _build_inputs( + batch_size, next_n, heads, head_dim, avg_kv_length, DTYPE_MAP[q_dtype] + ) + + with torch.inference_mode(): + ref = ref_fp8_paged_mqa_logits( + q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype + ) + ref_mask = ref == float("-inf") + + out = torch.full( + (batch_size * next_n, max_model_len), + float("-inf"), + device="cuda", + dtype=torch.float32, + ) + + def fn(): + out.fill_(float("-inf")) + return flydsl_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + out, + context_lens, + block_tables, + max_model_len, + SplitKV=_split_kv, + ) + + with torch.inference_mode(): + got, us = run_perftest(fn) + + got_mask = got == float("-inf") + assert torch.equal(got_mask, ref_mask), "flydsl paged: -inf mask mismatch" + + err = 0.0 + if not ref_mask.all(): + diff = calc_diff(got.masked_fill(got_mask, 0), ref.masked_fill(ref_mask, 0)) + assert diff < 1e-3, f"flydsl paged calc_diff={diff}" + err = diff.item() + checkAllclose( + ref.masked_fill(ref_mask, 0).to(dtypes.fp32), + got.masked_fill(got_mask, 0).to(dtypes.fp32), + rtol=1e-2, + atol=5.0, + msg="flydsl paged fp8_mqa_logits", + printLog=False, + ) + + total_flops = ( + 2 * next_n * heads * head_dim * context_lens.float().sum().item() + ) + return { + "gfx": get_gfx(), + "split_kv": "auto" if split_kv == 0 else split_kv, + "flydsl us": us, + "flydsl TFLOPS": total_flops / us / 1e6 if us > 0 else 0, + "flydsl err": err, + } + + +def main(): + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning( + "fp8_paged_mqa_logits unsupported on %s; skipping", get_gfx() + ) + return + if flydsl_fp8_paged_mqa_logits is None: + aiter.logger.warning("flydsl package not installed; skipping") + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="FlyDSL paged fp8_mqa_logits correctness + perf sweep", + ) + parser.add_argument( + "-b", + "--batch-next-n", + type=dtypes.str2tuple, + nargs="*", + default=[(1, 1), (1, 2), (2, 1), (2, 2), (4, 2), (8, 1)], + help="(batch_size, next_n) pairs", + ) + parser.add_argument("--num-heads", type=int, nargs="*", default=[64, 128]) + parser.add_argument("--head-dim", type=int, nargs="*", default=[64, 128]) + parser.add_argument( + "--avg-kv-length", type=int, nargs="*", default=[128, 1024, 8192] + ) + parser.add_argument( + "--q-dtype", + type=str, + nargs="*", + default=["fnuz", "fn"], + choices=["fnuz", "fn"], + ) + parser.add_argument( + "--split-kv", + type=int, + nargs="*", + default=[0, 1, 4], + help="0 == auto (production formula); else explicit SplitKV (1 disables)", + ) + args = parser.parse_args() + + df = [] + for (bs, nn), nh, hd, kv, qd, sk in itertools.product( + args.batch_next_n, + args.num_heads, + args.head_dim, + args.avg_kv_length, + args.q_dtype, + args.split_kv, + ): + df.append(test_fp8_paged_mqa_logits(bs, nn, nh, hd, kv, qd, sk)) + + df = pd.DataFrame(df) + try: + summary = df.to_markdown(index=False) + except ImportError: + # `tabulate` (optional) not installed -> plain-text table. + summary = df.to_string(index=False) + aiter.logger.info("fp8_paged_mqa_logits summary:\n%s", summary) + + +if __name__ == "__main__": + main() diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py new file mode 100644 index 00000000000..a585df76d19 --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py @@ -0,0 +1,315 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Honest A/B: FlyDSL paged MQA-logits vs the production Gluon/Triton kernel. + +Both kernels are timed by the SAME code path: identical inputs (built once and +fed to both), identical ``run_perftest`` methodology (device self-time, IQR +cleaning, same warmup/iters), identical output dtype/layout, and both are +correctness-gated against the SAME torch reference before any timing is reported. + +The reference kernel is the production ``deepgemm_fp8_paged_mqa_logits`` host +(aiter Triton/Gluon), run with vLLM's production defaults (ChunkK=256, +WavePerEU=2, KVBlockSize=1, Preshuffle=False) -- NOT detuned. Both the reference +host and the FlyDSL host auto-compute SplitKV to fill the device on small decode +grids, so the comparison is apples-to-apples. + +Reuses the paged benchmark's fp8 cache packer to guarantee the same input build. +""" + +import argparse +import random +import statistics + +import torch + +from aiter.test_common import run_perftest +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits +from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits + +# same fp8 co-pack builder the production paged benchmark uses +from op_tests.op_benchmarks.triton.bench_deepgemm_attention import ( + kv_cache_cast_to_fp8, + cdiv, +) + +torch.set_default_device("cuda") + +# Production reference config (matches vLLM rocm_fp8_paged_mqa_logits for decode). +REF_CHUNK_K = 256 +REF_WAVE_PER_EU = 2 + + +def calc_diff(x, y): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + return (1 - 2 * (x * y).sum() / denominator).item() + + +def ref_fp8_paged_mqa_logits( + q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype +): + """Vectorized torch reference (dequantizes the co-packed fp8 cache).""" + batch_size, next_n, heads, dim = q.size() + kvv, scale = kv_cache_fp8[..., :dim], kv_cache_fp8[..., dim:] + scale = scale.contiguous().view(torch.float) + qf = q.float() + kvf = (kvv.view(fp8_dtype).float() * scale).view(kv_cache_fp8.shape[0], dim) + logits = torch.full( + [batch_size * next_n, max_model_len], + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + for i in range(batch_size): + context_len = int(context_lens[i].item()) + if context_len == 0: + continue + pages = block_tables[i, :context_len] + kx = kvf[pages] + s = torch.einsum("nhd,pd->nhp", qf[i], kx) + s = torch.relu(s) + wl = weights[i * next_n : (i + 1) * next_n, :] + s = (s * wl[:, :, None]).sum(dim=1) + p = torch.arange(context_len, device=q.device) + q_lim = ( + context_len - next_n + torch.arange(next_n, device=q.device) + ).unsqueeze(1) + s = torch.where(p[None, :] <= q_lim, s, float("-inf")) + logits[i * next_n : (i + 1) * next_n, :context_len] = s + return logits + + +def parity_mask(context_lens, batch_size, next_n, max_model_len): + """Causal-valid mask (bench methodology): positions <= ctx - next_n + off.""" + positions = ( + torch.arange(max_model_len, device="cuda") + .unsqueeze(0) + .expand(batch_size * next_n, -1) + ) + row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n + next_n_offset = torch.arange(batch_size * next_n, device="cuda") % next_n + return positions <= ( + context_lens[row_indices] - next_n + next_n_offset + ).unsqueeze(1) + + +def build_inputs(batch_size, next_n, heads, index_dim, avg_kv_length, seed=0): + """Mirror the production paged benchmark's input builder (blocksize==1).""" + torch.manual_seed(seed) + random.seed(seed) + fp8_dtype = get_fp8_e4m3_dtype() + + max_model_len = 2 * avg_kv_length + blocksize = 1 + num_blocks = (max_model_len + blocksize - 1) // blocksize + + var_ratio = 0.5 + context_lens = ( + torch.randint( + int((1 - var_ratio) * avg_kv_length), + int((1 + var_ratio) * avg_kv_length) + 1, + (batch_size,), + ) + .cuda() + .to(torch.int32) + ) + # decode with MTP needs at least next_n tokens of context. + context_lens = torch.clamp(context_lens, min=next_n) + + q = torch.randn( + (batch_size, next_n, heads, index_dim), device="cuda", dtype=torch.bfloat16 + ) + kv_cache = torch.randn( + (num_blocks, blocksize, 1, index_dim), device="cuda", dtype=torch.bfloat16 + ) + weights = torch.randn( + (batch_size * next_n, heads), device="cuda", dtype=torch.float32 + ) + + max_block_len = (context_lens.max().item() + blocksize - 1) // blocksize * blocksize + block_tables = torch.zeros( + (batch_size, max_block_len), device="cuda", dtype=torch.int32 + ) + pool = list(range(num_blocks)) + random.shuffle(pool) + pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) + counter = 0 + for i in range(batch_size): + ctx_len = int(context_lens[i].item()) + n = cdiv(ctx_len, blocksize) + idx = (counter + torch.arange(n, device="cuda")) % num_blocks + block_tables[i, :n] = pool_t[idx] + counter += n + + q_fp8 = q.to(fp8_dtype) + kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, padding=False, fp8_dtype=fp8_dtype) + return ( + q, + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + ) + + +def time_kernel(fn, repeats, num_iters): + us = [] + for _ in range(repeats): + _, u = run_perftest(fn, num_iters=num_iters) + us.append(float(u)) + return statistics.median(us), min(us), max(us) + + +def run_shape(batch_size, next_n, heads, index_dim, avg_kv_length, repeats, num_iters): + ( + q, + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + ) = build_inputs(batch_size, next_n, heads, index_dim, avg_kv_length) + + ref = ref_fp8_paged_mqa_logits( + q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype + ) + valid = parity_mask(context_lens, batch_size, next_n, max_model_len) + neg_inf = float("-inf") + ref_neg_mask = ref == neg_inf + + def correctness(out): + diff = calc_diff(out.masked_fill(~valid, 0), ref.masked_fill(~valid, 0)) + mask_ok = bool(torch.equal(out == neg_inf, ref_neg_mask)) + return diff, mask_ok + + # ---- reference (production Gluon/Triton), production defaults ---- + out_ref = torch.full( + (batch_size * next_n, max_model_len), neg_inf, device="cuda", dtype=torch.float32 + ) + + def fn_ref(): + deepgemm_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + out_ref, + context_lens, + block_tables, + max_model_len, + ChunkK=REF_CHUNK_K, + Preshuffle=False, + KVBlockSize=1, + WavePerEU=REF_WAVE_PER_EU, + ) + + # ---- FlyDSL (KVBlockSize==1, auto SplitKV) ---- + out_fly = torch.full( + (batch_size * next_n, max_model_len), neg_inf, device="cuda", dtype=torch.float32 + ) + + def fn_fly(): + flydsl_fp8_paged_mqa_logits( + q_fp8, + kv_cache_fp8, + weights, + out_fly, + context_lens, + block_tables, + max_model_len, + ) + + # Warm / JIT-compile once (excluded from the reported device self-time anyway). + fn_ref() + fn_fly() + torch.cuda.synchronize() + + ref_diff, ref_mask_ok = correctness(out_ref) + fly_diff, fly_mask_ok = correctness(out_fly) + ref_pass = (ref_diff < 1e-3) and ref_mask_ok + fly_pass = (fly_diff < 1e-3) and fly_mask_ok + + ref_med = ref_min = ref_max = fly_med = fly_min = fly_max = float("nan") + if ref_pass and fly_pass: + ref_med, ref_min, ref_max = time_kernel(fn_ref, repeats, num_iters) + fly_med, fly_min, fly_max = time_kernel(fn_fly, repeats, num_iters) + + return { + "B": batch_size, + "nn": next_n, + "grid": batch_size * next_n, + "H": heads, + "D": index_dim, + "avg_kv": avg_kv_length, + "ref_ms": ref_med / 1e3, + "ref_min_ms": ref_min / 1e3, + "ref_max_ms": ref_max / 1e3, + "fly_ms": fly_med / 1e3, + "fly_min_ms": fly_min / 1e3, + "fly_max_ms": fly_max / 1e3, + "fly/ref": (fly_med / ref_med) if ref_med == ref_med and ref_med > 0 else float("nan"), + "ref_diff": ref_diff, + "fly_diff": fly_diff, + "ref_mask": ref_mask_ok, + "fly_mask": fly_mask_ok, + "ref_pass": ref_pass, + "fly_pass": fly_pass, + } + + +# DSA lightning-indexer decode dims (heads=64, index_dim=128) swept over long +# contexts and a concurrency (batch) range; next_n=2 is speculative decode, +# next_n=1 non-speculative. All shapes reported, wins and losses. +def default_shapes(): + H, D = 64, 128 + shapes = [] + for B in (1, 4, 16, 64, 128): + for avg_kv in (16384, 32768, 65536): + shapes.append((B, 2, H, D, avg_kv)) + for B in (1, 16, 128): + shapes.append((B, 1, H, D, 32768)) + return shapes + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--num-iters", type=int, default=101) + args = parser.parse_args() + + if get_gfx() not in ("gfx942", "gfx950"): + print(f"unsupported gfx {get_gfx()}; skipping") + return + + print(f"# arch={get_gfx()} fp8={get_fp8_e4m3_dtype()} " + f"repeats={args.repeats} num_iters={args.num_iters} " + f"ref_chunk_k={REF_CHUNK_K} ref_wave_per_eu={REF_WAVE_PER_EU}") + header = ( + "B nn grid H D avg_kv | ref_ms[min-max] fly_ms[min-max] fly/ref | " + "ref_diff fly_diff ref_mask fly_mask ref_pass fly_pass" + ) + print(header) + rows = [] + for (B, nn, H, D, avg) in default_shapes(): + r = run_shape(B, nn, H, D, avg, args.repeats, args.num_iters) + rows.append(r) + print( + f"{r['B']} {r['nn']} {r['grid']} {r['H']} {r['D']} {r['avg_kv']} | " + f"{r['ref_ms']:.4f}[{r['ref_min_ms']:.4f}-{r['ref_max_ms']:.4f}] " + f"{r['fly_ms']:.4f}[{r['fly_min_ms']:.4f}-{r['fly_max_ms']:.4f}] " + f"{r['fly/ref']:.2f}x | " + f"{r['ref_diff']:.2e} {r['fly_diff']:.2e} " + f"{r['ref_mask']} {r['fly_mask']} {r['ref_pass']} {r['fly_pass']}", + flush=True, + ) + + +if __name__ == "__main__": + main() From 65cba2b201e3cb4620fc164bde7df6a547b6d981 Mon Sep 17 00:00:00 2001 From: yashagar Date: Fri, 10 Jul 2026 09:40:45 +0000 Subject: [PATCH 02/52] fix bench to not use random avg context length --- .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 62 +++++++++++++------ 1 file changed, 43 insertions(+), 19 deletions(-) diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py index a585df76d19..43ce53918fa 100644 --- a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py @@ -96,26 +96,35 @@ def parity_mask(context_lens, batch_size, next_n, max_model_len): ).unsqueeze(1) -def build_inputs(batch_size, next_n, heads, index_dim, avg_kv_length, seed=0): - """Mirror the production paged benchmark's input builder (blocksize==1).""" +def build_inputs(batch_size, next_n, heads, index_dim, kv_length, seed=0, var_ratio=0.0): + """Paged input builder (blocksize==1). + + ``var_ratio`` controls the per-sequence context-length spread around + ``kv_length``: 0.0 gives an exact length (every sequence == kv_length), and + v>0 draws lengths uniformly from [(1-v)*kv_length, (1+v)*kv_length]. + """ torch.manual_seed(seed) random.seed(seed) fp8_dtype = get_fp8_e4m3_dtype() - max_model_len = 2 * avg_kv_length + max_model_len = 2 * kv_length blocksize = 1 num_blocks = (max_model_len + blocksize - 1) // blocksize - var_ratio = 0.5 - context_lens = ( - torch.randint( - int((1 - var_ratio) * avg_kv_length), - int((1 + var_ratio) * avg_kv_length) + 1, - (batch_size,), + if var_ratio == 0.0: + context_lens = torch.full( + (batch_size,), kv_length, device="cuda", dtype=torch.int32 + ) + else: + context_lens = ( + torch.randint( + int((1 - var_ratio) * kv_length), + int((1 + var_ratio) * kv_length) + 1, + (batch_size,), + ) + .cuda() + .to(torch.int32) ) - .cuda() - .to(torch.int32) - ) # decode with MTP needs at least next_n tokens of context. context_lens = torch.clamp(context_lens, min=next_n) @@ -166,7 +175,9 @@ def time_kernel(fn, repeats, num_iters): return statistics.median(us), min(us), max(us) -def run_shape(batch_size, next_n, heads, index_dim, avg_kv_length, repeats, num_iters): +def run_shape( + batch_size, next_n, heads, index_dim, kv_length, repeats, num_iters, var_ratio=0.0 +): ( q, q_fp8, @@ -176,7 +187,9 @@ def run_shape(batch_size, next_n, heads, index_dim, avg_kv_length, repeats, num_ block_tables, max_model_len, fp8_dtype, - ) = build_inputs(batch_size, next_n, heads, index_dim, avg_kv_length) + ) = build_inputs( + batch_size, next_n, heads, index_dim, kv_length, var_ratio=var_ratio + ) ref = ref_fp8_paged_mqa_logits( q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype @@ -247,7 +260,7 @@ def fn_fly(): "grid": batch_size * next_n, "H": heads, "D": index_dim, - "avg_kv": avg_kv_length, + "avg_kv": kv_length, "ref_ms": ref_med / 1e3, "ref_min_ms": ref_min / 1e3, "ref_max_ms": ref_max / 1e3, @@ -282,23 +295,34 @@ def main(): parser = argparse.ArgumentParser() parser.add_argument("--repeats", type=int, default=5) parser.add_argument("--num-iters", type=int, default=101) + parser.add_argument( + "--var-ratio", + type=float, + default=0.0, + help="Per-sequence context-length spread around kv_len. 0.0 = exact " + "(every sequence == kv_len); e.g. 0.5 draws lengths in [0.5, 1.5]*kv_len.", + ) args = parser.parse_args() if get_gfx() not in ("gfx942", "gfx950"): print(f"unsupported gfx {get_gfx()}; skipping") return + len_mode = "exact" if args.var_ratio == 0.0 else f"uniform +/-{args.var_ratio:g}" print(f"# arch={get_gfx()} fp8={get_fp8_e4m3_dtype()} " f"repeats={args.repeats} num_iters={args.num_iters} " - f"ref_chunk_k={REF_CHUNK_K} ref_wave_per_eu={REF_WAVE_PER_EU}") + f"ref_chunk_k={REF_CHUNK_K} ref_wave_per_eu={REF_WAVE_PER_EU} " + f"ctx_len={len_mode} max_model_len=2*kv_len blocksize=1") header = ( - "B nn grid H D avg_kv | ref_ms[min-max] fly_ms[min-max] fly/ref | " + "B nn grid H D kv_len | ref_ms[min-max] fly_ms[min-max] fly/ref | " "ref_diff fly_diff ref_mask fly_mask ref_pass fly_pass" ) print(header) rows = [] - for (B, nn, H, D, avg) in default_shapes(): - r = run_shape(B, nn, H, D, avg, args.repeats, args.num_iters) + for (B, nn, H, D, kv_len) in default_shapes(): + r = run_shape( + B, nn, H, D, kv_len, args.repeats, args.num_iters, var_ratio=args.var_ratio + ) rows.append(r) print( f"{r['B']} {r['nn']} {r['grid']} {r['H']} {r['D']} {r['avg_kv']} | " From 485f08b157a7fa1164ece46e693fe54d9724e601 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 10 Jul 2026 11:54:22 +0000 Subject: [PATCH 03/52] simplify fp8_paged_mqa_logits: operator style for i32 arithmetic, flatten variant scaffolding --- .../flydsl/kernels/fp8_paged_mqa_logits.py | 132 ++++++------------ 1 file changed, 41 insertions(+), 91 deletions(-) diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py index 92e93e6d03f..95943f6df50 100644 --- a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py @@ -51,7 +51,6 @@ MFMA_K, DEFAULT_COMPILE_HINTS, Vec, - i32_add as _i32_add, device_cu_count, load_pack_i64, fn_to_fnuz_i64, @@ -116,7 +115,6 @@ def _build_paged_kernel( N_TILES_PER_WAVE = N_TILES // WPB fm_fast = arith.FastMathFlags.fast - D_over_4 = D // 4 # f32-element offset of the co-packed scale within a token _cvt_tag = "" if convert_q_fn: @@ -161,7 +159,7 @@ def kernel( lane = fx.Int32(arith.remui(_to_raw(tid), _to_raw(fx.Int32(64)))) lane_div_N = fx.Int32(arith.divui(_to_raw(lane), _to_raw(fx.Int32(MFMA_N)))) lane_mod_N = fx.Int32(arith.remui(_to_raw(lane), _to_raw(fx.Int32(MFMA_N)))) - lane8 = fx.Int32(arith.muli(_to_raw(lane_div_N), _to_raw(fx.Int32(8)))) + lane8 = lane_div_N * 8 # fp8 operands read 8 bytes at a time as 2 i32 dwords (v8i8 buffer_load # fails to lower on gfx942), then bitcast to i64 for the MFMA. @@ -179,50 +177,29 @@ def kernel( context_length = fx.Int32(cl_t[pid_batch]) # Inclusive causal upper bound: p <= context_length - next_n + pid_next_n. - q_limit = fx.Int32( - arith.addi( - arith.subi(_to_raw(context_length), _to_raw(next_n)), - _to_raw(pid_next_n), - ) - ) + q_limit = context_length - next_n + pid_next_n - out_row = _i32_add( - fx.Int32(arith.muli(_to_raw(pid_batch), _to_raw(next_n))), pid_next_n - ) + out_row = pid_batch * next_n + pid_next_n stride_out_i64 = arith.extui(T.i64, _to_raw(stride_out)) out_row_t = make_out_row_view(out_logits, stride_out_i64, out_row) # ---- Preload Q frags + weights for the single query row ---- # A-operand layout is per in-wave lane, so `lane` (not `tid`) indexes Q. - q_row_base = _i32_add( - fx.Int32(arith.muli(_to_raw(pid_batch), _to_raw(stride_q_batch))), - fx.Int32(arith.muli(_to_raw(pid_next_n), _to_raw(stride_q_next_n))), - ) + q_row_base = pid_batch * stride_q_batch + pid_next_n * stride_q_next_n a_pack = [[None] * K_STEPS for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): - h_a = _i32_add(fx.Int32(mi * MFMA_M), lane_mod_N) - base_a = _i32_add( - q_row_base, - fx.Int32(arith.muli(_to_raw(h_a), _to_raw(stride_q_heads))), - ) + h_a = mi * MFMA_M + lane_mod_N + base_a = q_row_base + h_a * stride_q_heads for kk in range_constexpr(K_STEPS): - d_a = _i32_add(fx.Int32(kk * MFMA_K), lane8) - raw = load_pack_i64(q_i32, _i32_add(base_a, d_a)) + d_a = kk * MFMA_K + lane8 + raw = load_pack_i64(q_i32, base_a + d_a) a_pack[mi][kk] = fn_to_fnuz_i64(raw) if convert_q_fn else raw # weights[out_row, h] per (mi, ii): head = mi*MFMA_M + lane_div_N*4 + ii w_frag = [[None] * 4 for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): for ii in range_constexpr(4): - h_w = _i32_add( - fx.Int32(mi * MFMA_M), - _i32_add( - fx.Int32( - arith.muli(_to_raw(lane_div_N), _to_raw(fx.Int32(4))) - ), - fx.Int32(ii), - ), - ) + h_w = mi * MFMA_M + lane_div_N * 4 + ii w_frag[mi][ii] = _to_raw(fx.Float32(w_t[out_row, h_w])) # ---- SplitKV: this CTA owns a disjoint, BKV-aligned slice of the KV @@ -236,55 +213,37 @@ def kernel( _to_raw(context_length), _to_raw(fx.Int32(BKV)) ) split_chunk_num = arith.ceildivui(context_chunk_num, _to_raw(split_kv)) - full_end = arith.muli(context_chunk_num, _to_raw(fx.Int32(BKV))) - split_cols = arith.muli(split_chunk_num, _to_raw(fx.Int32(BKV))) - tile_lo_col = arith.muli(_to_raw(pid_split_kv), split_cols) - tile_hi_col = arith.minsi(arith.addi(tile_lo_col, split_cols), full_end) - ctx_m1 = arith.subi(_to_raw(context_length), _to_raw(fx.Int32(1))) + full_end = context_chunk_num * BKV + split_cols = split_chunk_num * BKV + tile_lo_col = pid_split_kv * split_cols + tile_hi_col = arith.minsi(_to_raw(tile_lo_col + split_cols), _to_raw(full_end)) + ctx_m1 = context_length - 1 - tile_lo = _to_raw(fx.Index(fx.Int32(tile_lo_col))) + tile_lo = _to_raw(fx.Index(tile_lo_col)) tile_hi = _to_raw(fx.Index(fx.Int32(tile_hi_col))) tile_step = _to_raw(fx.Index(fx.Int32(BKV))) tile_loop = scf.ForOp(tile_lo, tile_hi, tile_step, []) with ir.InsertionPoint(tile_loop.body): col0 = fx.Int32(arith.index_cast(T.i32, tile_loop.induction_variable)) - wave_ni_base = fx.Int32( - arith.muli(_to_raw(wave), _to_raw(fx.Int32(N_TILES_PER_WAVE))) - ) + wave_ni_base = wave * N_TILES_PER_WAVE for ni in range_constexpr(N_TILES_PER_WAVE): - abs_ni = _i32_add(wave_ni_base, fx.Int32(ni)) - col = _i32_add( - _i32_add( - col0, - fx.Int32( - arith.muli(_to_raw(abs_ni), _to_raw(fx.Int32(MFMA_N))) - ), - ), - lane_mod_N, - ) + abs_ni = wave_ni_base + ni + col = col0 + abs_ni * MFMA_N + lane_mod_N # Clamp the gather index to a valid token (mask handles the rest). - col_c = fx.Int32(arith.minsi(_to_raw(col), ctx_m1)) - ind_off = _i32_add( - fx.Int32(arith.muli(_to_raw(pid_batch), _to_raw(max_block_len))), - col_c, - ) + col_c = fx.Int32(arith.minsi(_to_raw(col), _to_raw(ctx_m1))) + ind_off = pid_batch * max_block_len + col_c physical = fx.Int32(ind_t[ind_off]) - tok_byte = fx.Int32( - arith.muli(_to_raw(physical), _to_raw(index_dim)) - ) + tok_byte = physical * index_dim b_col = [None] * K_STEPS for kk in range_constexpr(K_STEPS): - d_b = _i32_add(fx.Int32(kk * MFMA_K), lane8) - raw = load_pack_i64(kv_i32, _i32_add(tok_byte, d_b)) + d_b = kk * MFMA_K + lane8 + raw = load_pack_i64(kv_i32, tok_byte + d_b) b_col[kk] = fn_to_fnuz_i64(raw) if convert_kv_fn else raw # Co-packed per-token scale: f32 at byte tok_byte + D. scale_dword = fx.Int32( - arith.divui( - _to_raw(_i32_add(tok_byte, fx.Int32(D))), - _to_raw(fx.Int32(4)), - ) + arith.divui(_to_raw(tok_byte + D), _to_raw(fx.Int32(4))) ) kv_scale = _to_raw(fx.Float32(kv_f32[scale_dword])) if scale_mul != 1.0: @@ -376,33 +335,28 @@ def launch_fp8_paged_mqa_logits( # Kernel variants: single-token-per-block, wave-split only ("paged_w"). # WPB must divide the column-tile count BKV/16 (=8 at the default BKV=128). -def _mk_builder(wpb): - return lambda **kw: _build_paged_kernel(**kw, waves_per_block=wpb) - - -_VARIANT_BUILDERS = {f"paged_w{w}": _mk_builder(w) for w in (1, 2, 4)} -KERNEL_VARIANTS = tuple(_VARIANT_BUILDERS.keys()) +KERNEL_VARIANTS = tuple(f"paged_w{w}" for w in (1, 2, 4)) DEFAULT_VARIANT = "paged_w4" -def _auto_variant(batch_size, next_n, heads): - """Default variant selector: constant WPB=4 (8 column tiles / 4 waves = 2 - tiles/wave). Shape-adaptive variant selection is not implemented.""" - return DEFAULT_VARIANT +def _variant_wpb(variant): + """Parse the WPB int out of a ``paged_w`` tag (single validation point).""" + if variant not in KERNEL_VARIANTS: + raise ValueError( + f"unknown fp8_paged_mqa_logits variant {variant!r}; " + f"available: {list(KERNEL_VARIANTS)}" + ) + return int(variant.removeprefix("paged_w")) -def _resolve_variant(variant, batch_size, next_n, heads): - tag = ( +def _resolve_variant(variant): + """Pick the variant tag: explicit arg, then env override, then default. + Shape-adaptive selection is not implemented.""" + return ( variant or os.environ.get("FLYDSL_FP8_PAGED_MQA_LOGITS_VARIANT") - or _auto_variant(batch_size, next_n, heads) + or DEFAULT_VARIANT ) - if tag not in _VARIANT_BUILDERS: - raise ValueError( - f"unknown fp8_paged_mqa_logits variant {tag!r}; " - f"available: {list(KERNEL_VARIANTS)}" - ) - return tag @lru_cache(maxsize=32) @@ -423,15 +377,11 @@ def compile_fp8_paged_mqa_logits( operand whose -0 (0x80) byte the kernel patches to FNUZ +0, and ``scale_mul`` folds the 2x-per-converted-operand compensation into the co-packed scale. """ - if variant not in _VARIANT_BUILDERS: - raise ValueError( - f"unknown fp8_paged_mqa_logits variant {variant!r}; " - f"available: {list(KERNEL_VARIANTS)}" - ) - launcher = _VARIANT_BUILDERS[variant]( + launcher = _build_paged_kernel( num_heads=num_heads, head_size=head_size, block_kv=block_kv, + waves_per_block=_variant_wpb(variant), scale_mul=scale_mul, convert_q_fn=convert_q_fn, convert_kv_fn=convert_kv_fn, @@ -542,7 +492,7 @@ def flydsl_fp8_paged_mqa_logits( convert_kv_fn = False scale_mul = 2.0 if convert_q_fn else 1.0 - variant = _resolve_variant(variant, batch_size, next_n, num_heads) + variant = _resolve_variant(variant) launcher = compile_fp8_paged_mqa_logits( num_heads=num_heads, From b0bf341f638e069b27a61fc5e515cbaa22b0fa64 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 10 Jul 2026 13:38:52 +0000 Subject: [PATCH 04/52] =?UTF-8?q?perf:=20upgrade=20paged=20MQA=20logits=20?= =?UTF-8?q?to=20mfma=5Fscale=5Ff32=5F32x32x64=20(K=5FSTEPS=208=E2=86=922,?= =?UTF-8?q?=206-13%=20faster)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ops/flydsl/kernels/_mqa_logits_common.py | 75 +++++++++++++++---- aiter/ops/flydsl/kernels/fp8_mqa_logits.py | 13 ++-- .../flydsl/kernels/fp8_paged_mqa_logits.py | 38 ++++++---- 3 files changed, 91 insertions(+), 35 deletions(-) diff --git a/aiter/ops/flydsl/kernels/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/_mqa_logits_common.py index d6a49419594..147c4a2780c 100644 --- a/aiter/ops/flydsl/kernels/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/_mqa_logits_common.py @@ -32,16 +32,30 @@ Vec = fx.Vector -# MFMA tile dims of the fp8 16x16x32 atom: MFMA_M x MFMA_N output tile, -# MFMA_K fp8 elements reduced per MFMA step. -MFMA_M = 16 -MFMA_N = 16 -MFMA_K = 32 - -# The fp8 16x16x32 MFMA intrinsic. On gfx942 it reads operands as e4m3 FNUZ -# (bias 8); on gfx950 it reads the native e4m3 FN (OCP). See the FN->FNUZ patch -# below and the per-kernel ``convert_*_fn`` gates. -MFMA_FN = rocdl.mfma_f32_16x16x32_fp8_fp8 +# MFMA tile dims: 32x32x64 scaled fp8 (CDNA4, gfx950). +# K=64 per step → K_STEPS = head_size // 64 (2 for D=128). +# A/B operands are vector<8xi32> (4 packed i64s = 32 fp8 bytes per lane). +# 16 f32 D-regs/lane, same output tile shape as 32x32x16 (same D-reg layout). +# lane_div_N = lane // 32 (0 or 1), lane_mod_N = lane % 32 (KV col index). +# Scale operands use neutral e8m0 (0x7F7F7F7F = 1.0 per byte, 4 packed). +MFMA_M = 32 +MFMA_N = 32 +MFMA_K = 64 + +# The scaled 32x32x64 MFMA: A/B are vector<8xi32>, same D-reg layout as 32x32x16. +# operands: [a128, b128, c, cbsz=0, blgp=0, opsel_a=0, scale_a, opsel_b=0, scale_b] +# where a128/b128 are Vec(8xi32) (4 i64s packed), scale_a/scale_b are i32 values +# with 4 packed e8m0 bytes. Neutral: 0x7F7F7F7F (all 1.0). +def _mfma_scale_f32_32x32x64_fp8(result_type, operands): + from flydsl.expr.rocdl import _unwrap_mfma_operand + a = _unwrap_mfma_operand(operands[0]) # vector<8xi32> + b = _unwrap_mfma_operand(operands[1]) # vector<8xi32> + c = _unwrap_mfma_operand(operands[2]) # vector<16xf32> + # Neutral scale: 0x7F7F7F7F = 4 packed e8m0 bytes, each 127 → 2^(127-127)=1.0 + neutral = arith.constant(0x7F7F7F7F, type=T.i32) + return rocdl.mfma_scale_f32_32x32x64_f8f6f4_(result_type, a, b, c, 0, 0, 0, neutral, 0, neutral) + +MFMA_FN = _mfma_scale_f32_32x32x64_fp8 DEFAULT_COMPILE_HINTS = { "waves_per_eu": 2, @@ -76,6 +90,30 @@ def load_pack_i64(i32_view, byte_off_i32): return Vec(v2).bitcast(fx.Int64)[0].ir_value() +def load_pack_v8i32(i32_view, byte_off_i32, lane8): + """Load 32 fp8 bytes as vector<8xi32> for the 32x32x64 scaled MFMA. + + The 32x32x64 MFMA (K=64) A/B operands are vector<8xi32> per lane. + Each lane holds 32 fp8 bytes covering 4 consecutive K-groups of 16. + Same layout as 4 consecutive 32x32x16 K-steps: the per-lane K-bytes are + at positions [lane_div_N*8 + kk*16 + 0..7] for kk in 0..3. + + ``byte_off_i32`` is the base byte offset (tok_byte or q_row_base + h_a*stride) + for the first K-group. ``lane8 = lane_div_N * 8`` is the within-K-group offset. + Returns a vector<8xi32> value (MLIR Value). + """ + def _load_i64(off): + dw = fx.Int32(arith.divui(_to_raw(byte_off_i32 + lane8 + off), _to_raw(fx.Int32(4)))) + v2 = i32_view.vec_load((dw,), vec_size=2) + return Vec(v2).bitcast(fx.Int64)[0].ir_value() + + i64_0 = _load_i64(0) # K-group 0: k = lane_div_N*8 + 0..7 + i64_1 = _load_i64(16) # K-group 1: k = lane_div_N*8 + 16..23 + i64_2 = _load_i64(32) # K-group 2: k = lane_div_N*8 + 32..39 + i64_3 = _load_i64(48) # K-group 3: k = lane_div_N*8 + 48..55 + return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) + + def fn_to_fnuz_i64(raw_i64): """Map FN byte 0x80 (neg-zero) -> 0x00 in 8 packed fp8 bytes. @@ -147,32 +185,37 @@ def mfma_head_reduce( f32_0, fm_fast, mfma_fn=MFMA_FN, + dreg_count=16, ): """One column's logit: fp8 MFMA over heads, ReLU * weight sum, kv-scale, head reduce. - ``a_row`` : ``a_row[mi][kk]`` packed-i64 Q fragments for this query row. - ``b_col`` : ``b_col[kk]`` packed-i64 K fragments for this KV column tile. + ``a_row`` : ``a_row[mi][kk]`` Q fragments (i64 or vector<8xi32>) for this row. + ``b_col`` : ``b_col[kk]`` K fragments (same type as a_row) for this KV tile. ``w_row`` : ``w_row[mi][ii]`` f32 head weights for this query row. ``kv_scale``: f32 per-KV-token dequant scale (>= 0), hoisted out of the head sum (ReLU is positive-homogeneous, so ReLU(s*x)=s*ReLU(x) and the whole column sum is scaled once instead of per head). + ``dreg_count``: number of f32 D-regs per lane (16 for 32x32 tiles). Returns the reduced f32 ``col_sum`` (raw ir value). The in-wave (width 64) - head reduce uses ``shuffle_xor`` offsets 16 and 32. + head reduce uses shuffle_xor(32) for the 32x32 MFMA (lane_div_N in {0,1}). """ col_sum = _to_raw(f32_0) for mi in range_constexpr(m_tiles): - acc = Vec.filled(4, 0.0, fx.Float32) + acc = Vec.filled(dreg_count, 0.0, fx.Float32) for kk in range_constexpr(k_steps): acc = mfma_fn(res_ty, [a_row[mi][kk], b_col[kk], acc, 0, 0, 0]) - for ii in range_constexpr(4): + for ii in range_constexpr(dreg_count): score = Vec(acc)[ii].ir_value() relu = arith.maximumf(score, _to_raw(f32_0)) wsc = arith.MulFOp(relu, w_row[mi][ii], fastmath=fm_fast).result col_sum = arith.AddFOp(col_sum, wsc, fastmath=fm_fast).result col_sum = arith.MulFOp(col_sum, kv_scale, fastmath=fm_fast).result - for sh in (16, 32): + # For 32x32 tile: lane_div_N = lane // 32 (0 or 1); shuffle_xor(32) sums them. + # For 16x16 tile: lane_div_N = lane // 16 (0..3); shuffle_xor(16) + xor(32). + shuffles = (32,) if dreg_count == 16 else (16, 32) + for sh in shuffles: peer = _to_raw(ArithValue(col_sum).shuffle_xor(sh, 64)) col_sum = arith.AddFOp(col_sum, peer, fastmath=fm_fast).result return col_sum diff --git a/aiter/ops/flydsl/kernels/fp8_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_mqa_logits.py index ae2725c3373..93c3a1c4f34 100644 --- a/aiter/ops/flydsl/kernels/fp8_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_mqa_logits.py @@ -150,7 +150,8 @@ def kernel( num_splits: fx.Int32, # grid.y KV-column splits (1 == no split) ): f32_0 = arith.constant(0.0, type=T.f32) - mfma_res_ty = Vec.make_type(4, fx.Float32) + _DREG = 16 # D-regs per lane for 32x32x16 fp8 MFMA + mfma_res_ty = Vec.make_type(_DREG, fx.Float32) tid = fx.thread_idx.x bid = fx.block_idx.x @@ -213,17 +214,18 @@ def kernel( row_a[mi][kk] = fn_to_fnuz_i64(raw) if convert_q_fn else raw a_packs[j] = row_a - # weights[row, h] per (mi, ii): head = mi*MFMA_M + lane_div_N*4 + ii - row_w = [[None] * 4 for _ in range_constexpr(M_TILES)] + # weights[row, h] per (mi, ii): head = mi*MFMA_M + h_off(lane_div_N, ii) + # For 32x32x16 fp8: h_off = (ii%4) + lane_div_N*4 + (ii//4)*8 + row_w = [[None] * _DREG for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): - for ii in range_constexpr(4): + for ii in range_constexpr(_DREG): h_w = _i32_add( fx.Int32(mi * MFMA_M), _i32_add( fx.Int32( arith.muli(_to_raw(lane_div_N), _to_raw(fx.Int32(4))) ), - fx.Int32(ii), + fx.Int32((ii % 4) + (ii // 4) * 8), ), ) row_w[mi][ii] = _to_raw(fx.Float32(w_t[row, h_w])) @@ -317,6 +319,7 @@ def kernel( res_ty=mfma_res_ty, f32_0=f32_0, fm_fast=fm_fast, + dreg_count=_DREG, ) # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py index 95943f6df50..8c6fe6e6e06 100644 --- a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py @@ -53,6 +53,7 @@ Vec, device_cu_count, load_pack_i64, + load_pack_v8i32, fn_to_fnuz_i64, make_out_row_view, mfma_head_reduce, @@ -142,7 +143,7 @@ def kernel( stride_out: fx.Int32, # out_logits.stride(0) == max_model_len ): f32_0 = arith.constant(0.0, type=T.f32) - mfma_res_ty = Vec.make_type(4, fx.Float32) + mfma_res_ty = Vec.make_type(16, fx.Float32) scale_mul_c = arith.constant(float(scale_mul), type=T.f32) tid = fx.thread_idx.x @@ -161,8 +162,8 @@ def kernel( lane_mod_N = fx.Int32(arith.remui(_to_raw(lane), _to_raw(fx.Int32(MFMA_N)))) lane8 = lane_div_N * 8 - # fp8 operands read 8 bytes at a time as 2 i32 dwords (v8i8 buffer_load - # fails to lower on gfx942), then bitcast to i64 for the MFMA. + # fp8 operands loaded as vector<8xi32> per K-step (32 bytes = 4 i64s) + # for the 32x32x64 scaled MFMA. One vec_load_dwordx2 per i64 sub-step. q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) # Uniform-base views over the co-packed cache: per-token gather is done # via the (per-lane) byte offset, NOT the buffer base -- a per-lane base @@ -185,21 +186,28 @@ def kernel( # ---- Preload Q frags + weights for the single query row ---- # A-operand layout is per in-wave lane, so `lane` (not `tid`) indexes Q. + # With MFMA_K=64, each K-step loads a vector<8xi32> (4 packed i64s covering + # k = lane8 + [0,16,32,48]). fn_to_fnuz not supported for v8i32 operands; + # gfx950 uses native fp8 so convert_q_fn is always False on this arch. q_row_base = pid_batch * stride_q_batch + pid_next_n * stride_q_next_n a_pack = [[None] * K_STEPS for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): h_a = mi * MFMA_M + lane_mod_N base_a = q_row_base + h_a * stride_q_heads for kk in range_constexpr(K_STEPS): - d_a = kk * MFMA_K + lane8 - raw = load_pack_i64(q_i32, base_a + d_a) - a_pack[mi][kk] = fn_to_fnuz_i64(raw) if convert_q_fn else raw - - # weights[out_row, h] per (mi, ii): head = mi*MFMA_M + lane_div_N*4 + ii - w_frag = [[None] * 4 for _ in range_constexpr(M_TILES)] + byte_base = base_a + kk * MFMA_K + a_pack[mi][kk] = load_pack_v8i32(q_i32, byte_base, lane8) + + # weights[out_row, h] per (mi, ii): head = mi*MFMA_M + h_off(lane_div_N, ii) + # For 32x32x16 fp8 MFMA, the D-reg ii at lane_div_N encodes head: + # h_off = (ii % 4) + lane_div_N * 4 + (ii // 4) * 8 + # This is the same 4-interleave as 16x16x32 but extended to 16 D-regs. + _DREG = 16 + w_frag = [[None] * _DREG for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): - for ii in range_constexpr(4): - h_w = mi * MFMA_M + lane_div_N * 4 + ii + for ii in range_constexpr(_DREG): + h_off = (ii % 4) + lane_div_N * 4 + (ii // 4) * 8 + h_w = mi * MFMA_M + h_off w_frag[mi][ii] = _to_raw(fx.Float32(w_t[out_row, h_w])) # ---- SplitKV: this CTA owns a disjoint, BKV-aligned slice of the KV @@ -235,11 +243,12 @@ def kernel( physical = fx.Int32(ind_t[ind_off]) tok_byte = physical * index_dim + # B operands: vector<8xi32> per K-step (MFMA_K=64). + # convert_kv_fn is always False (gfx950 native fp8). b_col = [None] * K_STEPS for kk in range_constexpr(K_STEPS): - d_b = kk * MFMA_K + lane8 - raw = load_pack_i64(kv_i32, tok_byte + d_b) - b_col[kk] = fn_to_fnuz_i64(raw) if convert_kv_fn else raw + byte_base_b = tok_byte + kk * MFMA_K + b_col[kk] = load_pack_v8i32(kv_i32, byte_base_b, lane8) # Co-packed per-token scale: f32 at byte tok_byte + D. scale_dword = fx.Int32( @@ -261,6 +270,7 @@ def kernel( res_ty=mfma_res_ty, f32_0=f32_0, fm_fast=fm_fast, + dreg_count=_DREG, ) # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. From 48d887ebbb5bbeca5f7126c682453145656784da Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Mon, 13 Jul 2026 14:59:25 +0000 Subject: [PATCH 05/52] review: un-hardcode ChunkK / block_kv in paged mqa logits host Address PR review feedback (ROCm/aiter#4136): ChunkK was hardcoded to _BLOCK_KV in the public flydsl_fp8_paged_mqa_logits call despite compile_fp8_paged_mqa_logits already accepting block_kv. Now ChunkK is forwarded, defaulting to _BLOCK_KV (128) and validated to be a multiple of MFMA_N before dispatch. --- aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py index 8c6fe6e6e06..1132b816df6 100644 --- a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py @@ -411,7 +411,7 @@ def flydsl_fp8_paged_mqa_logits( *, Preshuffle=False, KVBlockSize=1, - ChunkK=256, + ChunkK=_BLOCK_KV, SplitKV=None, WavePerEU=2, TotalCuCount=None, @@ -440,8 +440,8 @@ def flydsl_fp8_paged_mqa_logits( own disjoint, gap-free BKV-aligned column ranges (one writer per column, no cross-CTA reduction). WavePerEU/TotalCuCount: inputs to the auto-SplitKV formula. - ChunkK: accepted for Triton-contract parity; ignored here (the FlyDSL - kernel tiles the KV window at a fixed ``block_kv`` width). + ChunkK: KV tile width (must be a multiple of MFMA_N=32). Mirrors + the Triton ChunkK parameter. Defaults to 128. Returns the same ``out_logits`` tensor (written in place). """ @@ -454,6 +454,10 @@ def flydsl_fp8_paged_mqa_logits( "FlyDSL paged fp8_mqa_logits supports KVBlockSize==1 only " f"(got KVBlockSize={KVBlockSize})." ) + if ChunkK % MFMA_N != 0: + raise ValueError( + f"ChunkK={ChunkK} must be a multiple of MFMA_N={MFMA_N}." + ) batch_size, next_n, num_heads, head_size = q_fp8.shape assert num_heads & (num_heads - 1) == 0, "num q. heads should be power of 2." @@ -507,7 +511,7 @@ def flydsl_fp8_paged_mqa_logits( launcher = compile_fp8_paged_mqa_logits( num_heads=num_heads, head_size=head_size, - block_kv=_BLOCK_KV, + block_kv=ChunkK, variant=variant, scale_mul=scale_mul, convert_q_fn=convert_q_fn, From afa18862e45c4bccb515485add5206e8ec998992 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 14 Jul 2026 11:17:34 +0000 Subject: [PATCH 06/52] style: apply ruff + black to flydsl mqa-logits files --- .../ops/flydsl/kernels/_mqa_logits_common.py | 30 +++++++++------- aiter/ops/flydsl/kernels/fp8_mqa_logits.py | 1 - .../flydsl/kernels/fp8_paged_mqa_logits.py | 35 ++++++++---------- .../test_flydsl_fp8_paged_mqa_logits.py | 16 ++++++--- .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 36 ++++++++++++------- 5 files changed, 68 insertions(+), 50 deletions(-) diff --git a/aiter/ops/flydsl/kernels/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/_mqa_logits_common.py index 147c4a2780c..a29703c840f 100644 --- a/aiter/ops/flydsl/kernels/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/_mqa_logits_common.py @@ -42,18 +42,23 @@ MFMA_N = 32 MFMA_K = 64 + # The scaled 32x32x64 MFMA: A/B are vector<8xi32>, same D-reg layout as 32x32x16. # operands: [a128, b128, c, cbsz=0, blgp=0, opsel_a=0, scale_a, opsel_b=0, scale_b] # where a128/b128 are Vec(8xi32) (4 i64s packed), scale_a/scale_b are i32 values # with 4 packed e8m0 bytes. Neutral: 0x7F7F7F7F (all 1.0). def _mfma_scale_f32_32x32x64_fp8(result_type, operands): from flydsl.expr.rocdl import _unwrap_mfma_operand - a = _unwrap_mfma_operand(operands[0]) # vector<8xi32> - b = _unwrap_mfma_operand(operands[1]) # vector<8xi32> - c = _unwrap_mfma_operand(operands[2]) # vector<16xf32> + + a = _unwrap_mfma_operand(operands[0]) # vector<8xi32> + b = _unwrap_mfma_operand(operands[1]) # vector<8xi32> + c = _unwrap_mfma_operand(operands[2]) # vector<16xf32> # Neutral scale: 0x7F7F7F7F = 4 packed e8m0 bytes, each 127 → 2^(127-127)=1.0 neutral = arith.constant(0x7F7F7F7F, type=T.i32) - return rocdl.mfma_scale_f32_32x32x64_f8f6f4_(result_type, a, b, c, 0, 0, 0, neutral, 0, neutral) + return rocdl.mfma_scale_f32_32x32x64_f8f6f4_( + result_type, a, b, c, 0, 0, 0, neutral, 0, neutral + ) + MFMA_FN = _mfma_scale_f32_32x32x64_fp8 @@ -102,15 +107,18 @@ def load_pack_v8i32(i32_view, byte_off_i32, lane8): for the first K-group. ``lane8 = lane_div_N * 8`` is the within-K-group offset. Returns a vector<8xi32> value (MLIR Value). """ + def _load_i64(off): - dw = fx.Int32(arith.divui(_to_raw(byte_off_i32 + lane8 + off), _to_raw(fx.Int32(4)))) + dw = fx.Int32( + arith.divui(_to_raw(byte_off_i32 + lane8 + off), _to_raw(fx.Int32(4))) + ) v2 = i32_view.vec_load((dw,), vec_size=2) return Vec(v2).bitcast(fx.Int64)[0].ir_value() - i64_0 = _load_i64(0) # K-group 0: k = lane_div_N*8 + 0..7 - i64_1 = _load_i64(16) # K-group 1: k = lane_div_N*8 + 16..23 - i64_2 = _load_i64(32) # K-group 2: k = lane_div_N*8 + 32..39 - i64_3 = _load_i64(48) # K-group 3: k = lane_div_N*8 + 48..55 + i64_0 = _load_i64(0) # K-group 0: k = lane_div_N*8 + 0..7 + i64_1 = _load_i64(16) # K-group 1: k = lane_div_N*8 + 16..23 + i64_2 = _load_i64(32) # K-group 2: k = lane_div_N*8 + 32..39 + i64_3 = _load_i64(48) # K-group 3: k = lane_div_N*8 + 48..55 return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) @@ -166,9 +174,7 @@ def make_out_row_view(logits, stride_i64, row_i32): mis-writing. """ _ri64 = arith.extui(T.i64, _to_raw(row_i32)) - _byte = arith.muli( - arith.muli(_ri64, stride_i64), arith.constant(4, type=T.i64) - ) + _byte = arith.muli(arith.muli(_ri64, stride_i64), arith.constant(4, type=T.i64)) _idx = arith.index_cast(T.index, _byte) return GTensor(logits, dtype=T.f32, shape=(-1,), static_bytes_offset_i64=_idx) diff --git a/aiter/ops/flydsl/kernels/fp8_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_mqa_logits.py index 93c3a1c4f34..b2c9b2cdbf6 100644 --- a/aiter/ops/flydsl/kernels/fp8_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_mqa_logits.py @@ -48,7 +48,6 @@ mfma_head_reduce, ) - # Default KV tile width (columns processed per inner-loop iteration). _BLOCK_KV = 128 # Don't split a row's KV window into chunks smaller than this many BKV tiles -- diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py index 1132b816df6..8d1a4ab3c42 100644 --- a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py @@ -52,13 +52,12 @@ DEFAULT_COMPILE_HINTS, Vec, device_cu_count, - load_pack_i64, load_pack_v8i32, - fn_to_fnuz_i64, make_out_row_view, mfma_head_reduce, ) + def _auto_split_kv(batch_size, next_n, wave_per_eu, device_index, total_cu=None): """Production host SplitKV formula (mirrors deepgemm_fp8_paged_mqa_logits): @@ -73,6 +72,7 @@ def _auto_split_kv(batch_size, next_n, wave_per_eu, device_index, total_cu=None) split_kv = ((max(1, total_cu // tile_q_count) + 4) // 5 * 5) * wave_per_eu return max(1, int(split_kv)) + # Default KV tile width (columns processed per MFMA inner-loop iteration). _BLOCK_KV = 128 @@ -336,9 +336,7 @@ def launch_fp8_paged_mqa_logits( index_dim, max_block_len, stride_out, - ).launch( - grid=(gx, 1, 1), block=(MR_BLOCK_THREADS, 1, 1), stream=stream - ) + ).launch(grid=(gx, 1, 1), block=(MR_BLOCK_THREADS, 1, 1), stream=stream) return launch_fp8_paged_mqa_logits @@ -455,25 +453,21 @@ def flydsl_fp8_paged_mqa_logits( f"(got KVBlockSize={KVBlockSize})." ) if ChunkK % MFMA_N != 0: - raise ValueError( - f"ChunkK={ChunkK} must be a multiple of MFMA_N={MFMA_N}." - ) + raise ValueError(f"ChunkK={ChunkK} must be a multiple of MFMA_N={MFMA_N}.") batch_size, next_n, num_heads, head_size = q_fp8.shape assert num_heads & (num_heads - 1) == 0, "num q. heads should be power of 2." assert head_size & (head_size - 1) == 0, "head size should be power of 2." num_blocks, block_size, one, index_dim = kv_cache.shape - assert block_size == 1, ( - f"kv_cache KVBlockSize dim must be 1; got {block_size}." - ) + assert block_size == 1, f"kv_cache KVBlockSize dim must be 1; got {block_size}." assert one == 1, f"kv_cache head dim must be 1; got {one}." - assert index_dim >= head_size + 4, ( - f"index_dim={index_dim} must hold {head_size} fp8 bytes + 4 scale bytes." - ) - assert kv_cache.dtype == torch.uint8, ( - f"kv_cache must be uint8 co-packed bytes; got {kv_cache.dtype}." - ) + assert ( + index_dim >= head_size + 4 + ), f"index_dim={index_dim} must hold {head_size} fp8 bytes + 4 scale bytes." + assert ( + kv_cache.dtype == torch.uint8 + ), f"kv_cache must be uint8 co-packed bytes; got {kv_cache.dtype}." # i32 gather-offset ceiling: the per-token byte base physical*index_dim is # computed in i32 (it rides the buffer voffset), so the whole cache pool must @@ -493,9 +487,10 @@ def flydsl_fp8_paged_mqa_logits( _fnuz = torch.float8_e4m3fnuz _fn = torch.float8_e4m3fn - assert q_fp8.dtype in (_fnuz, _fn), ( - f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" - ) + assert q_fp8.dtype in ( + _fnuz, + _fn, + ), f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" # gfx942 fp8 MFMA reads operands as e4m3 FNUZ (bias 8). Q may arrive as FN # (OCP, bias 7); patch its 0x80 byte and undo the implied 2x via scale_mul. diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 0a25a4d6da7..fd8e3e82b43 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -109,7 +109,9 @@ def ref_fp8_paged_mqa_logits( p = torch.arange(context_len, device=q.device) q_lim = ( context_len - next_n + torch.arange(next_n, device=q.device) - ).unsqueeze(1) # [next_n, 1] + ).unsqueeze( + 1 + ) # [next_n, 1] s = torch.where(p[None, :] <= q_lim, s, float("-inf")) logits[i * next_n : (i + 1) * next_n, :context_len] = s return logits @@ -181,7 +183,13 @@ def test_fp8_paged_mqa_logits( with torch.inference_mode(): ref = ref_fp8_paged_mqa_logits( - q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype + q, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, ) ref_mask = ref == float("-inf") @@ -225,9 +233,7 @@ def fn(): printLog=False, ) - total_flops = ( - 2 * next_n * heads * head_dim * context_lens.float().sum().item() - ) + total_flops = 2 * next_n * heads * head_dim * context_lens.float().sum().item() return { "gfx": get_gfx(), "split_kv": "auto" if split_kv == 0 else split_kv, diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py index 43ce53918fa..08f6a322554 100644 --- a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py @@ -91,12 +91,14 @@ def parity_mask(context_lens, batch_size, next_n, max_model_len): ) row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n next_n_offset = torch.arange(batch_size * next_n, device="cuda") % next_n - return positions <= ( - context_lens[row_indices] - next_n + next_n_offset - ).unsqueeze(1) + return positions <= (context_lens[row_indices] - next_n + next_n_offset).unsqueeze( + 1 + ) -def build_inputs(batch_size, next_n, heads, index_dim, kv_length, seed=0, var_ratio=0.0): +def build_inputs( + batch_size, next_n, heads, index_dim, kv_length, seed=0, var_ratio=0.0 +): """Paged input builder (blocksize==1). ``var_ratio`` controls the per-sequence context-length spread around @@ -205,7 +207,10 @@ def correctness(out): # ---- reference (production Gluon/Triton), production defaults ---- out_ref = torch.full( - (batch_size * next_n, max_model_len), neg_inf, device="cuda", dtype=torch.float32 + (batch_size * next_n, max_model_len), + neg_inf, + device="cuda", + dtype=torch.float32, ) def fn_ref(): @@ -225,7 +230,10 @@ def fn_ref(): # ---- FlyDSL (KVBlockSize==1, auto SplitKV) ---- out_fly = torch.full( - (batch_size * next_n, max_model_len), neg_inf, device="cuda", dtype=torch.float32 + (batch_size * next_n, max_model_len), + neg_inf, + device="cuda", + dtype=torch.float32, ) def fn_fly(): @@ -267,7 +275,9 @@ def fn_fly(): "fly_ms": fly_med / 1e3, "fly_min_ms": fly_min / 1e3, "fly_max_ms": fly_max / 1e3, - "fly/ref": (fly_med / ref_med) if ref_med == ref_med and ref_med > 0 else float("nan"), + "fly/ref": ( + (fly_med / ref_med) if ref_med == ref_med and ref_med > 0 else float("nan") + ), "ref_diff": ref_diff, "fly_diff": fly_diff, "ref_mask": ref_mask_ok, @@ -309,17 +319,19 @@ def main(): return len_mode = "exact" if args.var_ratio == 0.0 else f"uniform +/-{args.var_ratio:g}" - print(f"# arch={get_gfx()} fp8={get_fp8_e4m3_dtype()} " - f"repeats={args.repeats} num_iters={args.num_iters} " - f"ref_chunk_k={REF_CHUNK_K} ref_wave_per_eu={REF_WAVE_PER_EU} " - f"ctx_len={len_mode} max_model_len=2*kv_len blocksize=1") + print( + f"# arch={get_gfx()} fp8={get_fp8_e4m3_dtype()} " + f"repeats={args.repeats} num_iters={args.num_iters} " + f"ref_chunk_k={REF_CHUNK_K} ref_wave_per_eu={REF_WAVE_PER_EU} " + f"ctx_len={len_mode} max_model_len=2*kv_len blocksize=1" + ) header = ( "B nn grid H D kv_len | ref_ms[min-max] fly_ms[min-max] fly/ref | " "ref_diff fly_diff ref_mask fly_mask ref_pass fly_pass" ) print(header) rows = [] - for (B, nn, H, D, kv_len) in default_shapes(): + for B, nn, H, D, kv_len in default_shapes(): r = run_shape( B, nn, H, D, kv_len, args.repeats, args.num_iters, var_ratio=args.var_ratio ) From 50bb92d4129e49430d2ddb4516c5c1272d929d69 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Thu, 13 Aug 2026 13:06:45 +0000 Subject: [PATCH 07/52] feat(flydsl): paged mqa-logits KVBlockSize>=1 + correctness-only test Kernel: extend fp8_paged_mqa_logits from KVBlockSize==1 to KVBlockSize>=1 (block-flat co-packed layout: KVB fp8 key rows grouped first, then KVB f32 scales). Per-lane gather resolves each logical KV column independently through the block table (kv_indices[b, col // KVB] -> block, col % KVB -> token), so KVB is a compile-time constant that folds to the per-token slot layout at KVB==1. i32 gather-offset ceiling now accounts for KVB*index_dim. Test: add --kv-block-size sweep axis (1, 64) and reference block-flat dequant; convert to a pure correctness test -- drop run_perftest / the fn() wrapper for a direct kernel call, and remove the timing-derived us/TFLOPS output columns (gate stays exact -inf mask + calc_diff < 1e-3). Verified on gfx950 from a fresh compile cache across KVB {1,64}, dtypes {fnuz,fn}, SplitKV {auto,1,4}, heads/head_dim {64,128}, short/long context. --- .../flydsl/kernels/fp8_paged_mqa_logits.py | 126 ++++++++++++------ .../test_flydsl_fp8_paged_mqa_logits.py | 114 +++++++++++----- 2 files changed, 162 insertions(+), 78 deletions(-) diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py index 8d1a4ab3c42..318db64bd88 100644 --- a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py @@ -14,16 +14,23 @@ maps through the block table (``kv_indices``) to a physical block, and each token's fp8 K bytes and its f32 dequant scale are **co-packed** in that block. -Supported scope (this file): ``KVBlockSize == 1`` with SplitKV KV-column -parallelism. The public ``flydsl_fp8_paged_mqa_logits`` mirrors the tensor -contract of the Triton +Supported scope (this file): ``KVBlockSize >= 1`` (block-flat co-packed layout) +with SplitKV KV-column parallelism. The public ``flydsl_fp8_paged_mqa_logits`` +mirrors the tensor contract of the Triton ``aiter.ops.triton.attention.pa_mqa_logits.deepgemm_fp8_paged_mqa_logits`` so the -two are interchangeable in tests/benchmarks. ``KVBlockSize > 1`` and preshuffle -are not supported and raise ``NotImplementedError``. - -The fp8 16x16x32 MFMA compute, ReLU*weight head-sum + kv-scale hoist, in-wave -``shuffle_xor`` head reduce, the fp8 dword-pack load, the FN->FNUZ byte patch, and -the i64 byte-base per-row output view are all shared with the dense kernel via +two are interchangeable in tests/benchmarks. Preshuffle is not supported and +raises ``NotImplementedError``. + +Cache layout (per physical block of ``KVBlockSize`` tokens): all ``KVBlockSize`` +fp8 key rows (``D`` bytes each) grouped first, then all ``KVBlockSize`` f32 +dequant scales (block-flat). The per-lane gather resolves each logical KV column +independently through the block table (``kv_indices[b, p // KVBlockSize]``) so a +column tile may straddle block boundaries; ``KVBlockSize == 1`` degenerates to +the per-token ``[D fp8 | 4 scale]`` slot layout. + +The fp8 32x32x64 scaled MFMA compute, ReLU*weight head-sum + kv-scale hoist, +in-wave ``shuffle_xor`` head reduce, the fp8 dword-pack load, and the i64 +byte-base per-row output view are all shared with the dense kernel via ``._mqa_logits_common`` (no duplication). """ @@ -83,17 +90,25 @@ def _build_paged_kernel( head_size: int, block_kv: int, waves_per_block: int, + kv_block_size: int = 1, scale_mul: float = 1.0, convert_q_fn: bool = False, convert_kv_fn: bool = False, ): - """Paged MQA-logits kernel (KVBlockSize==1) with SplitKV column parallelism. + """Paged MQA-logits kernel (KVBlockSize>=1) with SplitKV column parallelism. One thread block owns one ``(batch, next_n)`` query row and walks the whole ``[0, context_length)`` KV window in ``BKV``-wide tiles. ``waves_per_block`` waves each own a disjoint slice of the ``BKV/MFMA_N`` column tiles (no cross-wave sharing / barrier), mirroring the dense kernel's wave split. + ``kv_block_size`` (KVBlockSize) is the number of tokens co-packed in one + physical block; each lane gathers its own logical KV column independently + (``kv_indices[b, col // KVBlockSize]`` -> physical block, ``col % + KVBlockSize`` -> token within the block-flat block), so a column tile may + straddle block boundaries. ``KVBlockSize == 1`` degenerates to the per-token + ``[D fp8 | 4 scale]`` slot layout. + ``scale_mul`` folds the FN->FNUZ 2x-per-converted-operand compensation into the (co-packed) per-token scale at compile time (ReLU positive-homogeneity). """ @@ -101,6 +116,8 @@ def _build_paged_kernel( D = head_size BKV = block_kv WPB = waves_per_block + KVB = kv_block_size + assert KVB >= 1, f"kv_block_size must be >= 1; got {KVB}" MR_BLOCK_THREADS = 64 * WPB assert H % MFMA_M == 0, f"num_heads={H} must be a multiple of MFMA_M={MFMA_M}" @@ -122,12 +139,12 @@ def _build_paged_kernel( _cvt_tag += "_cq" if convert_kv_fn: _cvt_tag += "_ck" - _kname = f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_w{WPB}{_cvt_tag}_flydsl" + _kname = f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_kvb{KVB}_w{WPB}{_cvt_tag}_flydsl" @flyc.kernel(name=_kname, known_block_size=[MR_BLOCK_THREADS, 1, 1]) def kernel( Q: fx.Tensor, # [batch, next_n, H, D] fp8 (bytes passed raw) - KV_cache: fx.Tensor, # [num_blocks, 1, 1, index_dim] uint8 co-packed + KV_cache: fx.Tensor, # [num_blocks, KVB, 1, index_dim] uint8 block-flat weights: fx.Tensor, # [batch*next_n, H] f32 out_logits: fx.Tensor, # [batch*next_n, max_model_len] f32 (-inf prefilled) context_lens: fx.Tensor, # [batch] i32 @@ -138,7 +155,7 @@ def kernel( stride_q_batch: fx.Int32, # fp8 elems (== bytes) stride_q_next_n: fx.Int32, stride_q_heads: fx.Int32, - index_dim: fx.Int32, # cache row width in bytes (D + 4 + pad) + index_dim: fx.Int32, # per-token slot bytes (D+4+pad); block stride = KVB*index_dim max_block_len: fx.Int32, # kv_indices row width stride_out: fx.Int32, # out_logits.stride(0) == max_model_len ): @@ -199,9 +216,9 @@ def kernel( a_pack[mi][kk] = load_pack_v8i32(q_i32, byte_base, lane8) # weights[out_row, h] per (mi, ii): head = mi*MFMA_M + h_off(lane_div_N, ii) - # For 32x32x16 fp8 MFMA, the D-reg ii at lane_div_N encodes head: - # h_off = (ii % 4) + lane_div_N * 4 + (ii // 4) * 8 - # This is the same 4-interleave as 16x16x32 but extended to 16 D-regs. + # For the 32x32x64 MFMA (16 f32 D-regs), the D-reg ii at lane_div_N encodes + # head: h_off = (ii % 4) + lane_div_N * 4 + (ii // 4) * 8 + # -- the same 4-interleave as the 16x16x32 layout, extended to 16 D-regs. _DREG = 16 w_frag = [[None] * _DREG for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): @@ -226,6 +243,9 @@ def kernel( tile_lo_col = pid_split_kv * split_cols tile_hi_col = arith.minsi(_to_raw(tile_lo_col + split_cols), _to_raw(full_end)) ctx_m1 = context_length - 1 + # Block-flat block byte-stride (KVB tokens per physical block); + # loop-invariant, so hoist it out of the per-column gather below. + block_byte_stride = fx.Int32(KVB) * index_dim tile_lo = _to_raw(fx.Index(tile_lo_col)) tile_hi = _to_raw(fx.Index(fx.Int32(tile_hi_col))) @@ -239,20 +259,35 @@ def kernel( col = col0 + abs_ni * MFMA_N + lane_mod_N # Clamp the gather index to a valid token (mask handles the rest). col_c = fx.Int32(arith.minsi(_to_raw(col), _to_raw(ctx_m1))) - ind_off = pid_batch * max_block_len + col_c + # Block table maps a logical column to a physical block. The + # block-flat block byte-stride is KVB*index_dim, keys grouped + # first (tok*D) then all scales (KVB*D + tok*4). KVB is a + # compile-time constant, so for KVB==1 divui/remui/mul-by-1 fold + # away (block_idx==col_c, tok_in_block==0) -- no Python `if` here + # because the FlyDSL AST rewriter would turn it into a runtime + # scf.if and drop the branch-local addresses. + block_idx = fx.Int32( + arith.divui(_to_raw(col_c), _to_raw(fx.Int32(KVB))) + ) + tok_in_block = fx.Int32( + arith.remui(_to_raw(col_c), _to_raw(fx.Int32(KVB))) + ) + ind_off = pid_batch * max_block_len + block_idx physical = fx.Int32(ind_t[ind_off]) - tok_byte = physical * index_dim + block_byte_base = physical * block_byte_stride + key_byte = block_byte_base + tok_in_block * D # B operands: vector<8xi32> per K-step (MFMA_K=64). # convert_kv_fn is always False (gfx950 native fp8). b_col = [None] * K_STEPS for kk in range_constexpr(K_STEPS): - byte_base_b = tok_byte + kk * MFMA_K + byte_base_b = key_byte + kk * MFMA_K b_col[kk] = load_pack_v8i32(kv_i32, byte_base_b, lane8) - # Co-packed per-token scale: f32 at byte tok_byte + D. + # Co-packed per-token f32 scale (block-flat: after the KVB keys). + scale_byte = block_byte_base + (KVB * D) + tok_in_block * 4 scale_dword = fx.Int32( - arith.divui(_to_raw(tok_byte + D), _to_raw(fx.Int32(4))) + arith.divui(_to_raw(scale_byte), _to_raw(fx.Int32(4))) ) kv_scale = _to_raw(fx.Float32(kv_f32[scale_dword])) if scale_mul != 1.0: @@ -373,6 +408,7 @@ def compile_fp8_paged_mqa_logits( num_heads: int, head_size: int, block_kv: int = _BLOCK_KV, + kv_block_size: int = 1, variant: str = DEFAULT_VARIANT, scale_mul: float = 1.0, convert_q_fn: bool = False, @@ -380,16 +416,18 @@ def compile_fp8_paged_mqa_logits( ): """Return a cached, compiled FlyDSL paged launcher for the given config. - ``num_heads``/``head_size`` are compile-time constants; ``variant`` is a - ``paged_w`` tag; ``convert_q_fn``/``convert_kv_fn`` mark an FP8 FN - operand whose -0 (0x80) byte the kernel patches to FNUZ +0, and ``scale_mul`` - folds the 2x-per-converted-operand compensation into the co-packed scale. + ``num_heads``/``head_size``/``kv_block_size`` are compile-time constants; + ``variant`` is a ``paged_w`` tag; ``convert_q_fn``/``convert_kv_fn`` mark + an FP8 FN operand whose -0 (0x80) byte the kernel patches to FNUZ +0, and + ``scale_mul`` folds the 2x-per-converted-operand compensation into the + co-packed scale. """ launcher = _build_paged_kernel( num_heads=num_heads, head_size=head_size, block_kv=block_kv, waves_per_block=_variant_wpb(variant), + kv_block_size=kv_block_size, scale_mul=scale_mul, convert_q_fn=convert_q_fn, convert_kv_fn=convert_kv_fn, @@ -416,21 +454,23 @@ def flydsl_fp8_paged_mqa_logits( variant=None, stream=None, ): - """FlyDSL paged FP8 MQA logits (decode) -- KVBlockSize==1 with SplitKV. + """FlyDSL paged FP8 MQA logits (decode) -- KVBlockSize>=1 with SplitKV. Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. q_fp8: [batch, next_n, heads, hidden_dim], dtype float8 (e4m3 fn/fnuz) - kv_cache: [num_blocks, KVBlockSize, 1, index_dim], dtype uint8, co-packed - fp8 K bytes (first KVBlockSize*hidden_dim) then f32 dequant - scales; index_dim == hidden_dim + 4 (+ optional 16B padding). + kv_cache: [num_blocks, KVBlockSize, 1, index_dim], dtype uint8, block-flat + co-packed: per physical block, KVBlockSize*hidden_dim fp8 K + bytes (all key rows grouped) then KVBlockSize f32 dequant scales; + index_dim == hidden_dim + 4 (+ optional 16B padding). weights: [batch*next_n, heads], dtype float32 out_logits: [batch*next_n, max_model_len], dtype float32. MUST be prefilled with -inf by the caller; the kernel writes only in-window (causal) positions and leaves masked positions untouched. context_lens: [batch], dtype int32 kv_indices: [batch, max_block_len], dtype int32 block table (physical block - per logical position; KVBlockSize==1 => per-token page table). + per logical block ``p // KVBlockSize``; KVBlockSize==1 => a + per-token page table). max_model_len: int, out_logits column count. SplitKV: KV-column split count (grid = split_kv*batch*next_n). None => the production host formula (fills the device on small decode @@ -447,11 +487,6 @@ def flydsl_fp8_paged_mqa_logits( raise NotImplementedError( "Preshuffle is not supported by the FlyDSL paged fp8_mqa_logits kernel." ) - if KVBlockSize != 1: - raise NotImplementedError( - "FlyDSL paged fp8_mqa_logits supports KVBlockSize==1 only " - f"(got KVBlockSize={KVBlockSize})." - ) if ChunkK % MFMA_N != 0: raise ValueError(f"ChunkK={ChunkK} must be a multiple of MFMA_N={MFMA_N}.") @@ -460,7 +495,11 @@ def flydsl_fp8_paged_mqa_logits( assert head_size & (head_size - 1) == 0, "head size should be power of 2." num_blocks, block_size, one, index_dim = kv_cache.shape - assert block_size == 1, f"kv_cache KVBlockSize dim must be 1; got {block_size}." + assert block_size == KVBlockSize, ( + f"kv_cache KVBlockSize dim ({block_size}) must equal KVBlockSize " + f"({KVBlockSize})." + ) + assert KVBlockSize >= 1, f"KVBlockSize must be >= 1; got {KVBlockSize}." assert one == 1, f"kv_cache head dim must be 1; got {one}." assert ( index_dim >= head_size + 4 @@ -469,11 +508,13 @@ def flydsl_fp8_paged_mqa_logits( kv_cache.dtype == torch.uint8 ), f"kv_cache must be uint8 co-packed bytes; got {kv_cache.dtype}." - # i32 gather-offset ceiling: the per-token byte base physical*index_dim is - # computed in i32 (it rides the buffer voffset), so the whole cache pool must - # be addressable in i32 bytes. A larger pool would need an i64/global gather. - assert num_blocks * index_dim < 2**31, ( - f"num_blocks*index_dim={num_blocks * index_dim} exceeds the i32 " + # i32 gather-offset ceiling: the per-token byte base (physical block byte base + # + intra-block offset) is computed in i32 (it rides the buffer voffset), so + # the whole cache pool must be addressable in i32 bytes. A larger pool would + # need an i64/global gather. + pool_bytes = num_blocks * block_size * index_dim + assert pool_bytes < 2**31, ( + f"num_blocks*KVBlockSize*index_dim={pool_bytes} exceeds the i32 " f"gather-offset limit (2^31); the paged kernel needs an i64 gather path " f"for a cache pool this large." ) @@ -507,13 +548,14 @@ def flydsl_fp8_paged_mqa_logits( num_heads=num_heads, head_size=head_size, block_kv=ChunkK, + kv_block_size=int(KVBlockSize), variant=variant, scale_mul=scale_mul, convert_q_fn=convert_q_fn, convert_kv_fn=convert_kv_fn, ) - # Co-packed cache -> raw uint8 byte view [num_blocks * index_dim]. + # Co-packed cache -> raw uint8 byte view [num_blocks * KVBlockSize * index_dim]. kv_bytes = kv_cache.reshape(-1) # KV-column splits (one lever): fill the device when the batch*next_n row diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index fd8e3e82b43..f1a454b7c54 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness + perf sweep for the FlyDSL paged FP8 MQA-logits (decode) kernel. +"""Correctness sweep for the FlyDSL paged FP8 MQA-logits (decode) kernel. Mirrors the aiter-op-test standard used by ``test_flydsl_fp8_mqa_logits.py`` and the paged benchmark ``op_tests/op_benchmarks/triton/bench_deepgemm_attention.py``. @@ -19,7 +19,7 @@ import pandas as pd import torch from aiter import dtypes -from aiter.test_common import benchmark, checkAllclose, run_perftest +from aiter.test_common import benchmark, checkAllclose from aiter.jit.utils.chip_info import get_gfx from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype @@ -72,24 +72,44 @@ def kv_cache_cast_to_fp8(x, fp8_dtype): def ref_fp8_paged_mqa_logits( - q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype + q, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + block_size=1, ): """Torch reference (vectorized port of vLLM ``fp8_paged_mqa_logits_torch``). - Dequantizes the co-packed fp8 cache (``fp8_bytes.float() * scale``) so it - matches what the kernel computes (kv-scale folded into K), then: + Dequantizes the **block-flat** co-packed fp8 cache and computes: logits[b*next_n+n, p] = sum_h ReLU() * weights[.., h] - with the causal mask ``p <= context_len - next_n + n``. The inner per-token - loop of the original is replaced by a single gather + einsum per batch - element (block_size == 1), which is orders of magnitude faster at long - context while numerically identical. + with the causal mask ``p <= context_len - next_n + n``. Logical position ``p`` + resolves to physical ``(block, tok) = (block_tables[b, p//KVB], p%KVB)`` and + the block-flat block groups all ``KVB`` fp8 key rows first, then all ``KVB`` + f32 scales. ``block_size == 1`` degenerates to the per-token page table. """ batch_size, next_n, heads, dim = q.size() - kvv, scale = kv_cache_fp8[..., :dim], kv_cache_fp8[..., dim:] - scale = scale.contiguous().view(torch.float) + num_blocks = kv_cache_fp8.shape[0] + index_dim = kv_cache_fp8.shape[-1] + # Reconstruct the block-flat block: [KVB keys (dim bytes each)][KVB f32 scales]. + flat = kv_cache_fp8.reshape(num_blocks, block_size * index_dim) + keys = ( + flat[:, : block_size * dim] + .contiguous() + .view(fp8_dtype) + .float() + .view(num_blocks, block_size, dim) + ) + scales = ( + flat[:, block_size * dim : block_size * dim + 4 * block_size] + .contiguous() + .view(torch.float32) + .view(num_blocks, block_size, 1) + ) + kvf = keys * scales # [num_blocks, block_size, dim] dequantized K qf = q.float() - # dequantized K per physical block: [num_blocks, dim] (block_size == 1) - kvf = (kvv.view(fp8_dtype).float() * scale).view(kv_cache_fp8.shape[0], dim) logits = torch.full( [batch_size * next_n, max_model_len], float("-inf"), @@ -100,30 +120,33 @@ def ref_fp8_paged_mqa_logits( context_len = int(context_lens[i].item()) if context_len == 0: continue - pages = block_tables[i, :context_len] # [ctx] physical token ids - kx = kvf[pages] # [ctx, dim] dequantized K + pos = torch.arange(context_len, device=q.device) + blk = block_tables[i, pos // block_size] # [ctx] physical block ids + tok = pos % block_size # [ctx] token-in-block + kx = kvf[blk, tok] # [ctx, dim] dequantized K s = torch.einsum("nhd,pd->nhp", qf[i], kx) # [next_n, heads, ctx] s = torch.relu(s) wl = weights[i * next_n : (i + 1) * next_n, :] # [next_n, heads] s = (s * wl[:, :, None]).sum(dim=1) # [next_n, ctx] - p = torch.arange(context_len, device=q.device) q_lim = ( context_len - next_n + torch.arange(next_n, device=q.device) ).unsqueeze( 1 ) # [next_n, 1] - s = torch.where(p[None, :] <= q_lim, s, float("-inf")) + s = torch.where(pos[None, :] <= q_lim, s, float("-inf")) logits[i * next_n : (i + 1) * next_n, :context_len] = s return logits -def _build_inputs(batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, seed=0): +def _build_inputs( + batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, block_size=1, seed=0 +): torch.manual_seed(seed) random.seed(seed) fp8_dtype = get_fp8_e4m3_dtype() max_model_len = 2 * avg_kv_length - num_blocks = max_model_len # KVBlockSize == 1 + num_blocks = (max_model_len + block_size - 1) // block_size lo = max(1, int((1 - 0.5) * avg_kv_length)) hi = int((1 + 0.5) * avg_kv_length) + 1 @@ -132,10 +155,11 @@ def _build_inputs(batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, s context_lens = torch.clamp(context_lens, min=next_n) q = torch.randn((batch_size, next_n, heads, head_dim), dtype=torch.bfloat16) - kv_cache = torch.randn((num_blocks, 1, 1, head_dim), dtype=torch.bfloat16) + kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) weights = torch.randn((batch_size * next_n, heads), dtype=torch.float32) - max_block_len = int(context_lens.max().item()) # blocksize == 1 + # block table: one entry per KVBlockSize-token block, ceil(ctx / block_size). + max_block_len = (int(context_lens.max().item()) + block_size - 1) // block_size block_tables = torch.zeros( (batch_size, max_block_len), device="cuda", dtype=torch.int32 ) @@ -143,7 +167,8 @@ def _build_inputs(batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, s random.shuffle(pool) counter = 0 for i in range(batch_size): - for j in range(int(context_lens[i].item())): + n_blk = (int(context_lens[i].item()) + block_size - 1) // block_size + for j in range(n_blk): block_tables[i][j] = pool[counter % num_blocks] counter += 1 @@ -163,7 +188,14 @@ def _build_inputs(batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, s @benchmark() def test_fp8_paged_mqa_logits( - batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, split_kv=0 + batch_size, + next_n, + heads, + head_dim, + avg_kv_length, + q_dtype, + split_kv=0, + block_size=1, ): # split_kv == 0 -> auto (production host formula); else an explicit override # (1 disables splitting). Both must be correctness-identical. @@ -178,7 +210,13 @@ def test_fp8_paged_mqa_logits( max_model_len, fp8_dtype, ) = _build_inputs( - batch_size, next_n, heads, head_dim, avg_kv_length, DTYPE_MAP[q_dtype] + batch_size, + next_n, + heads, + head_dim, + avg_kv_length, + DTYPE_MAP[q_dtype], + block_size=block_size, ) with torch.inference_mode(): @@ -190,6 +228,7 @@ def test_fp8_paged_mqa_logits( block_tables, max_model_len, fp8_dtype, + block_size=block_size, ) ref_mask = ref == float("-inf") @@ -200,9 +239,8 @@ def test_fp8_paged_mqa_logits( dtype=torch.float32, ) - def fn(): - out.fill_(float("-inf")) - return flydsl_fp8_paged_mqa_logits( + with torch.inference_mode(): + got = flydsl_fp8_paged_mqa_logits( q_fp8, kv_cache_fp8, weights, @@ -210,12 +248,10 @@ def fn(): context_lens, block_tables, max_model_len, + KVBlockSize=block_size, SplitKV=_split_kv, ) - with torch.inference_mode(): - got, us = run_perftest(fn) - got_mask = got == float("-inf") assert torch.equal(got_mask, ref_mask), "flydsl paged: -inf mask mismatch" @@ -233,12 +269,10 @@ def fn(): printLog=False, ) - total_flops = 2 * next_n * heads * head_dim * context_lens.float().sum().item() return { "gfx": get_gfx(), + "kvb": block_size, "split_kv": "auto" if split_kv == 0 else split_kv, - "flydsl us": us, - "flydsl TFLOPS": total_flops / us / 1e6 if us > 0 else 0, "flydsl err": err, } @@ -255,7 +289,7 @@ def main(): parser = argparse.ArgumentParser( formatter_class=argparse.RawTextHelpFormatter, - description="FlyDSL paged fp8_mqa_logits correctness + perf sweep", + description="FlyDSL paged fp8_mqa_logits correctness sweep", ) parser.add_argument( "-b", @@ -284,18 +318,26 @@ def main(): default=[0, 1, 4], help="0 == auto (production formula); else explicit SplitKV (1 disables)", ) + parser.add_argument( + "--kv-block-size", + type=int, + nargs="*", + default=[1, 64], + help="KVBlockSize (paged block-flat page size); 1 == per-token slots", + ) args = parser.parse_args() df = [] - for (bs, nn), nh, hd, kv, qd, sk in itertools.product( + for (bs, nn), nh, hd, kv, qd, sk, kvb in itertools.product( args.batch_next_n, args.num_heads, args.head_dim, args.avg_kv_length, args.q_dtype, args.split_kv, + args.kv_block_size, ): - df.append(test_fp8_paged_mqa_logits(bs, nn, nh, hd, kv, qd, sk)) + df.append(test_fp8_paged_mqa_logits(bs, nn, nh, hd, kv, qd, sk, block_size=kvb)) df = pd.DataFrame(df) try: From 4d70992269194c728c3b406963a6bc59deda348f Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Thu, 13 Aug 2026 13:49:56 +0000 Subject: [PATCH 08/52] feat(flydsl): add Preshuffle support to paged fp8 mqa-logits + isolation bench Support the production shuffle_weight(16x16) KV data layout in the FlyDSL paged fp8 MQA-logits kernel (Preshuffle=True, KVBlockSize % 16 == 0), making it a true drop-in for deepgemm_fp8_paged_mqa_logits(Preshuffle=True). Only the intra-block key-byte offset changes; the block-table gather, co-packed/token-ordered scale gather, MFMA, and causal mask are unchanged. - _mqa_logits_common: add load_pack_v8i32_preshuffle (closed-form offset into the shuffle_weight(16x16) data section; 8-byte halves stay contiguous). - fp8_paged_mqa_logits: thread a compile-time preshuffle flag through the kernel builder (with a _ps name tag), compile_* cache key, and the public API; replace the NotImplementedError guard with the KVBlockSize % 16 == 0 assert. - test: add a preshuffle path (oracle from the unshuffled cache, kernel fed the shuffled copy) and a small dedicated case set (kvb 16/64, D 64/128, FN patch). - bench_flydsl_vs_gluon: note the A/B is fixed at KVBlockSize=1 (Gluon faults at KVBlockSize>1); KVBlockSize>1 / Preshuffle coverage lives in the flydsl test. - add bench_flydsl_paged_mqa_logits.py: standalone single-config bench (no Gluon) for rocprofv3 isolation profiling, covering both block-flat and preshuffle paths. Validated on gfx950: block-flat and preshuffle correctness pass the unchanged gate (exact -inf mask + calc_diff < 1e-3); rocprofv3 traces the _ps kernel. --- .../ops/flydsl/kernels/_mqa_logits_common.py | 46 +++ .../flydsl/kernels/fp8_paged_mqa_logits.py | 64 +++- .../test_flydsl_fp8_paged_mqa_logits.py | 67 ++++- .../triton/bench_flydsl_paged_mqa_logits.py | 277 ++++++++++++++++++ .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 6 + 5 files changed, 443 insertions(+), 17 deletions(-) create mode 100644 op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py diff --git a/aiter/ops/flydsl/kernels/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/_mqa_logits_common.py index a29703c840f..5cfe447f452 100644 --- a/aiter/ops/flydsl/kernels/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/_mqa_logits_common.py @@ -122,6 +122,52 @@ def _load_i64(off): return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) +def load_pack_v8i32_preshuffle( + i32_view, block_byte_base, tok_in_block, head_size, kk_step, lane8 +): + """Load 32 fp8 bytes as vector<8xi32> from the production preshuffled layout. + + The production ``Preshuffle`` path stores each physical block's fp8 key data + reordered by ``shuffle_weight(layout=(16,16))`` (only the ``KVB*D`` data + section; the co-packed f32 scale tail is left token-ordered, so the scale + gather is unchanged). Element ``(token c, hidden h)`` lands at the intra-block + byte offset:: + + off(c,h) = (c//16)*(D*16) + (h//16)*256 + (c%16)*16 + (h%16) + + which is exactly the inverse of the host shuffle permute (and matches the + gluon CDNA ``offset_k_fixed``). This kernel's per-lane K payload for MFMA + K-step ``kk_step`` is the four contiguous hidden-16 groups + ``g = kk_step*4 + sub`` (``sub`` in 0..3); within a group the 8-byte half + picked by ``lane8`` (``= lane_div_N*8``) is still contiguous, so we issue the + same four ``dwordx2`` loads and assemble a ``vector<8xi32>``. + + ``block_byte_base`` is the physical block's byte base in the cache pool + (``physical * KVB * index_dim``); the shuffled data section starts at + offset 0 within it. ``head_size`` (``D``) must be a multiple of 32 (the 16x16 + shuffle pairs ``k_width=16`` fp8 into 32-byte chunks) -- guaranteed here by + ``D % MFMA_K == 0``. + """ + D = head_size + # Token super-row (16-token blocks) + intra-super-row column; loop-invariant + # across the four hidden groups, so hoist out of the per-group load. + c_hi = fx.Int32(arith.divui(_to_raw(tok_in_block), _to_raw(fx.Int32(16)))) + c_lo = fx.Int32(arith.remui(_to_raw(tok_in_block), _to_raw(fx.Int32(16)))) + base = block_byte_base + c_hi * (D * 16) + c_lo * 16 + lane8 + + def _load_i64(g): + dw = fx.Int32(arith.divui(_to_raw(base + g * 256), _to_raw(fx.Int32(4)))) + v2 = i32_view.vec_load((dw,), vec_size=2) + return Vec(v2).bitcast(fx.Int64)[0].ir_value() + + g0 = kk_step * 4 # first hidden-16 group of this MFMA K-step + i64_0 = _load_i64(g0 + 0) + i64_1 = _load_i64(g0 + 1) + i64_2 = _load_i64(g0 + 2) + i64_3 = _load_i64(g0 + 3) + return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) + + def fn_to_fnuz_i64(raw_i64): """Map FN byte 0x80 (neg-zero) -> 0x00 in 8 packed fp8 bytes. diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py index 318db64bd88..5e9e33d0b35 100644 --- a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py @@ -18,8 +18,9 @@ with SplitKV KV-column parallelism. The public ``flydsl_fp8_paged_mqa_logits`` mirrors the tensor contract of the Triton ``aiter.ops.triton.attention.pa_mqa_logits.deepgemm_fp8_paged_mqa_logits`` so the -two are interchangeable in tests/benchmarks. Preshuffle is not supported and -raises ``NotImplementedError``. +two are interchangeable in tests/benchmarks. ``Preshuffle=True`` consumes the +production ``shuffle_weight(layout=(16,16))`` KV data layout (requires +``KVBlockSize % 16 == 0``). Cache layout (per physical block of ``KVBlockSize`` tokens): all ``KVBlockSize`` fp8 key rows (``D`` bytes each) grouped first, then all ``KVBlockSize`` f32 @@ -60,6 +61,7 @@ Vec, device_cu_count, load_pack_v8i32, + load_pack_v8i32_preshuffle, make_out_row_view, mfma_head_reduce, ) @@ -91,6 +93,7 @@ def _build_paged_kernel( block_kv: int, waves_per_block: int, kv_block_size: int = 1, + preshuffle: bool = False, scale_mul: float = 1.0, convert_q_fn: bool = False, convert_kv_fn: bool = False, @@ -109,6 +112,12 @@ def _build_paged_kernel( straddle block boundaries. ``KVBlockSize == 1`` degenerates to the per-token ``[D fp8 | 4 scale]`` slot layout. + ``preshuffle`` selects the production ``shuffle_weight(layout=(16,16))`` KV + data layout (each block's ``KVB*D`` fp8 bytes reordered into 16-token x + 16-hidden tiles; the co-packed f32 scale tail stays token-ordered). Only the + intra-block key-byte offset changes -- the block-table gather, scale gather, + MFMA, and mask are identical. Requires ``KVBlockSize % 16 == 0``. + ``scale_mul`` folds the FN->FNUZ 2x-per-converted-operand compensation into the (co-packed) per-token scale at compile time (ReLU positive-homogeneity). """ @@ -139,7 +148,11 @@ def _build_paged_kernel( _cvt_tag += "_cq" if convert_kv_fn: _cvt_tag += "_ck" - _kname = f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_kvb{KVB}_w{WPB}{_cvt_tag}_flydsl" + _ps_tag = "_ps" if preshuffle else "" + _kname = ( + f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_kvb{KVB}_w{WPB}" + f"{_ps_tag}{_cvt_tag}_flydsl" + ) @flyc.kernel(name=_kname, known_block_size=[MR_BLOCK_THREADS, 1, 1]) def kernel( @@ -275,14 +288,24 @@ def kernel( ind_off = pid_batch * max_block_len + block_idx physical = fx.Int32(ind_t[ind_off]) block_byte_base = physical * block_byte_stride - key_byte = block_byte_base + tok_in_block * D # B operands: vector<8xi32> per K-step (MFMA_K=64). - # convert_kv_fn is always False (gfx950 native fp8). + # convert_kv_fn is always False (gfx950 native fp8). The + # preshuffle branch is compile-time (Python if): it only swaps + # the intra-block key-byte offset for the shuffle_weight(16x16) + # layout; the scale gather below is unchanged (scales stay + # co-packed and token-ordered). b_col = [None] * K_STEPS - for kk in range_constexpr(K_STEPS): - byte_base_b = key_byte + kk * MFMA_K - b_col[kk] = load_pack_v8i32(kv_i32, byte_base_b, lane8) + if preshuffle: + for kk in range_constexpr(K_STEPS): + b_col[kk] = load_pack_v8i32_preshuffle( + kv_i32, block_byte_base, tok_in_block, D, kk, lane8 + ) + else: + key_byte = block_byte_base + tok_in_block * D + for kk in range_constexpr(K_STEPS): + byte_base_b = key_byte + kk * MFMA_K + b_col[kk] = load_pack_v8i32(kv_i32, byte_base_b, lane8) # Co-packed per-token f32 scale (block-flat: after the KVB keys). scale_byte = block_byte_base + (KVB * D) + tok_in_block * 4 @@ -409,6 +432,7 @@ def compile_fp8_paged_mqa_logits( head_size: int, block_kv: int = _BLOCK_KV, kv_block_size: int = 1, + preshuffle: bool = False, variant: str = DEFAULT_VARIANT, scale_mul: float = 1.0, convert_q_fn: bool = False, @@ -416,11 +440,11 @@ def compile_fp8_paged_mqa_logits( ): """Return a cached, compiled FlyDSL paged launcher for the given config. - ``num_heads``/``head_size``/``kv_block_size`` are compile-time constants; - ``variant`` is a ``paged_w`` tag; ``convert_q_fn``/``convert_kv_fn`` mark - an FP8 FN operand whose -0 (0x80) byte the kernel patches to FNUZ +0, and - ``scale_mul`` folds the 2x-per-converted-operand compensation into the - co-packed scale. + ``num_heads``/``head_size``/``kv_block_size``/``preshuffle`` are compile-time + constants; ``variant`` is a ``paged_w`` tag; ``convert_q_fn``/ + ``convert_kv_fn`` mark an FP8 FN operand whose -0 (0x80) byte the kernel + patches to FNUZ +0, and ``scale_mul`` folds the 2x-per-converted-operand + compensation into the co-packed scale. """ launcher = _build_paged_kernel( num_heads=num_heads, @@ -428,6 +452,7 @@ def compile_fp8_paged_mqa_logits( block_kv=block_kv, waves_per_block=_variant_wpb(variant), kv_block_size=kv_block_size, + preshuffle=preshuffle, scale_mul=scale_mul, convert_q_fn=convert_q_fn, convert_kv_fn=convert_kv_fn, @@ -462,7 +487,12 @@ def flydsl_fp8_paged_mqa_logits( kv_cache: [num_blocks, KVBlockSize, 1, index_dim], dtype uint8, block-flat co-packed: per physical block, KVBlockSize*hidden_dim fp8 K bytes (all key rows grouped) then KVBlockSize f32 dequant scales; - index_dim == hidden_dim + 4 (+ optional 16B padding). + index_dim == hidden_dim + 4 (+ optional 16B padding). When + ``Preshuffle=True`` the fp8 data section is reordered by the + production ``shuffle_weight(layout=(16,16))`` (16-token x + 16-hidden tiles); the f32 scale tail stays token-ordered. + Preshuffle: consume the shuffle_weight(16x16) KV data layout. Requires + ``KVBlockSize % 16 == 0``. Default False (block-flat gather). weights: [batch*next_n, heads], dtype float32 out_logits: [batch*next_n, max_model_len], dtype float32. MUST be prefilled with -inf by the caller; the kernel writes only in-window @@ -484,8 +514,9 @@ def flydsl_fp8_paged_mqa_logits( Returns the same ``out_logits`` tensor (written in place). """ if Preshuffle: - raise NotImplementedError( - "Preshuffle is not supported by the FlyDSL paged fp8_mqa_logits kernel." + assert KVBlockSize % 16 == 0, ( + f"Preshuffle mode only supports KVBlockSize aligned to 16; " + f"got KVBlockSize={KVBlockSize}." ) if ChunkK % MFMA_N != 0: raise ValueError(f"ChunkK={ChunkK} must be a multiple of MFMA_N={MFMA_N}.") @@ -549,6 +580,7 @@ def flydsl_fp8_paged_mqa_logits( head_size=head_size, block_kv=ChunkK, kv_block_size=int(KVBlockSize), + preshuffle=bool(Preshuffle), variant=variant, scale_mul=scale_mul, convert_q_fn=convert_q_fn, diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index f1a454b7c54..2cd38c013d8 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -71,6 +71,33 @@ def kv_cache_cast_to_fp8(x, fp8_dtype): return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) +def preshuffle_kv_data(kv_cache_fp8, head_dim): + """Reorder the fp8 data section into the production ``Preshuffle`` layout. + + Applies ``shuffle_weight(layout=(16,16))`` to the per-block ``KVB*head_dim`` + fp8 key bytes (reshaped ``[num_blocks, KVBlockSize, head_dim]``) and leaves + the co-packed f32 scale tail untouched -- exactly what the production + ``deepgemm_fp8_paged_mqa_logits`` benchmark builds for ``Preshuffle=True``. + The torch oracle still reads the *unshuffled* cache (the reference is + layout-agnostic), so this shuffled copy is fed only to the kernel. + """ + from aiter.ops.shuffle import shuffle_weight + + num_blocks, block_size, one, index_dim = kv_cache_fp8.shape + assert block_size % 16 == 0, "preshuffle requires KVBlockSize % 16 == 0" + flat = kv_cache_fp8.reshape(num_blocks, block_size * index_dim).clone() + data = ( + flat[:, : block_size * head_dim] + .contiguous() + .view(num_blocks, block_size, head_dim) + ) + shuffled = shuffle_weight(data, layout=(16, 16)).reshape( + num_blocks, block_size * head_dim + ) + flat[:, : block_size * head_dim] = shuffled + return flat.view(num_blocks, block_size, one, index_dim) + + def ref_fp8_paged_mqa_logits( q, kv_cache_fp8, @@ -196,6 +223,7 @@ def test_fp8_paged_mqa_logits( q_dtype, split_kv=0, block_size=1, + preshuffle=False, ): # split_kv == 0 -> auto (production host formula); else an explicit override # (1 disables splitting). Both must be correctness-identical. @@ -239,15 +267,22 @@ def test_fp8_paged_mqa_logits( dtype=torch.float32, ) + # Oracle reads the unshuffled cache above; the kernel gets the preshuffled + # copy (same numbers, production shuffle_weight(16x16) data layout). + kv_cache_kernel = ( + preshuffle_kv_data(kv_cache_fp8, head_dim) if preshuffle else kv_cache_fp8 + ) + with torch.inference_mode(): got = flydsl_fp8_paged_mqa_logits( q_fp8, - kv_cache_fp8, + kv_cache_kernel, weights, out, context_lens, block_tables, max_model_len, + Preshuffle=preshuffle, KVBlockSize=block_size, SplitKV=_split_kv, ) @@ -272,11 +307,28 @@ def test_fp8_paged_mqa_logits( return { "gfx": get_gfx(), "kvb": block_size, + "preshuffle": preshuffle, "split_kv": "auto" if split_kv == 0 else split_kv, "flydsl err": err, } +# Preshuffle is a distinct correctness branch (production shuffle_weight(16x16) +# data layout, KVBlockSize % 16 == 0). A small fixed set covers what differs from +# the block-flat path -- KVBlockSize 16 and 64, head_dim 64 and 128, the FN->FNUZ +# Q patch, and a couple of grid/context shapes -- without re-running the whole +# block-flat sweep. +# (batch_size, next_n, heads, head_dim, avg_kv, q_dtype, kv_block_size) +PRESHUFFLE_CASES = [ + (1, 1, 64, 128, 1024, "fnuz", 16), + (1, 2, 64, 128, 1024, "fnuz", 64), + (2, 1, 128, 128, 8192, "fnuz", 64), + (1, 1, 64, 64, 1024, "fnuz", 16), # head_dim=64 (single K-step) + (1, 1, 64, 128, 1024, "fn", 64), # exercises the gfx942 FN->FNUZ Q patch + (4, 2, 64, 128, 8192, "fnuz", 64), +] + + def main(): if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning( @@ -325,6 +377,11 @@ def main(): default=[1, 64], help="KVBlockSize (paged block-flat page size); 1 == per-token slots", ) + parser.add_argument( + "--no-preshuffle", + action="store_true", + help="skip the preshuffle (shuffle_weight 16x16 layout) correctness cases", + ) args = parser.parse_args() df = [] @@ -339,6 +396,14 @@ def main(): ): df.append(test_fp8_paged_mqa_logits(bs, nn, nh, hd, kv, qd, sk, block_size=kvb)) + if not args.no_preshuffle: + for bs, nn, nh, hd, kv, qd, kvb in PRESHUFFLE_CASES: + df.append( + test_fp8_paged_mqa_logits( + bs, nn, nh, hd, kv, qd, 0, block_size=kvb, preshuffle=True + ) + ) + df = pd.DataFrame(df) try: summary = df.to_markdown(index=False) diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py new file mode 100644 index 00000000000..3593d2157fb --- /dev/null +++ b/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py @@ -0,0 +1,277 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Standalone microbenchmark for the FlyDSL paged FP8 MQA-logits (decode) kernel. + +Runs ONE config of ``flydsl_fp8_paged_mqa_logits`` in isolation (no reference / +Gluon comparison) so the kernel can be profiled directly. Example rocprofv3 run:: + + PYTHONPATH=. rocprofv3 --kernel-trace --output-format csv \\ + --output-directory .rocprofv3 -- \\ + python op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py \\ + --heads 64 --head-dim 128 --kv-len 32768 --batch 16 --next-n 2 \\ + --kv-block-size 64 --preshuffle --iters 200 --no-check + +Covers both the block-flat gather (``--kv-block-size N``, default 1) and the +production preshuffle layout (``--preshuffle``; requires KVBlockSize % 16 == 0). +The input build, co-pack byte layout, preshuffle, and torch oracle are imported +from the flydsl correctness test, so the profiled kernel sees exactly the tested +layout (single source of truth). + +By default the run is correctness-gated once (torch oracle, exact ``-inf`` mask + +calc_diff < 1e-3) before the timed loop; pass ``--no-check`` to skip it (handy +when profiling). The timed loop launches the kernel ``--iters`` times back to +back -- that is what rocprofv3 traces. ``--time`` additionally reports device +self-time via ``run_perftest``. +""" + +import argparse +import random + +import torch + +from aiter.test_common import run_perftest +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits + +# Reuse the correctness test's layout/oracle helpers so the profiled kernel sees +# exactly the tested co-pack + preshuffle byte layout. +from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( + kv_cache_cast_to_fp8, + ref_fp8_paged_mqa_logits, + preshuffle_kv_data, + calc_diff, +) + +torch.set_default_device("cuda") + +DTYPE_MAP = {"fnuz": get_fp8_e4m3_dtype(), "fn": torch.float8_e4m3fn} + + +def build_inputs( + batch, next_n, heads, head_dim, kv_len, block_size, q_dtype, var_ratio, seed +): + """Paged decode inputs (block-flat co-packed cache), block_size >= 1. + + ``var_ratio`` controls the per-sequence context-length spread around + ``kv_len``: 0.0 == exact (every sequence == kv_len); v>0 draws lengths in + [(1-v), (1+v)]*kv_len. A fixed ``seed`` makes the whole build deterministic + (stable kernel work across profiling runs). + """ + torch.manual_seed(seed) + random.seed(seed) + fp8_dtype = get_fp8_e4m3_dtype() + + max_model_len = 2 * kv_len + num_blocks = (max_model_len + block_size - 1) // block_size + + if var_ratio == 0.0: + context_lens = torch.full((batch,), kv_len, device="cuda", dtype=torch.int32) + else: + lo = max(1, int((1 - var_ratio) * kv_len)) + hi = int((1 + var_ratio) * kv_len) + 1 + context_lens = torch.randint(lo, hi, (batch,)).cuda().to(torch.int32) + # decode with MTP needs at least next_n tokens of context. + context_lens = torch.clamp(context_lens, min=next_n) + + q = torch.randn((batch, next_n, heads, head_dim), dtype=torch.bfloat16) + kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) + weights = torch.randn((batch * next_n, heads), dtype=torch.float32) + + # block table: ceil(ctx / block_size) physical blocks per sequence. + max_block_len = (int(context_lens.max().item()) + block_size - 1) // block_size + block_tables = torch.zeros( + (batch, max_block_len), device="cuda", dtype=torch.int32 + ) + pool = list(range(num_blocks)) + random.shuffle(pool) + counter = 0 + for i in range(batch): + n_blk = (int(context_lens[i].item()) + block_size - 1) // block_size + for j in range(n_blk): + block_tables[i][j] = pool[counter % num_blocks] + counter += 1 + + q_fp8 = q.to(q_dtype) + kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, fp8_dtype) + return ( + q, + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + ) + + +def main(): + parser = argparse.ArgumentParser( + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + description=__doc__, + ) + # Single-config by design: profile exactly one kernel launch shape. + parser.add_argument("--batch", type=int, default=16) + parser.add_argument("--next-n", type=int, default=2) + parser.add_argument("--heads", type=int, default=64) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--kv-len", type=int, default=32768) + parser.add_argument( + "--kv-block-size", type=int, default=1, help="KVBlockSize (block-flat page)" + ) + parser.add_argument( + "--preshuffle", + action="store_true", + help="use the shuffle_weight(16x16) KV layout (needs kv-block-size %% 16 == 0)", + ) + parser.add_argument( + "--q-dtype", type=str, default="fnuz", choices=["fnuz", "fn"] + ) + parser.add_argument( + "--split-kv", type=int, default=0, help="0 == auto (host formula); else override" + ) + parser.add_argument("--wave-per-eu", type=int, default=2) + parser.add_argument("--chunk-k", type=int, default=128, help="KV tile width") + parser.add_argument( + "--variant", type=str, default=None, help="paged_w tag (default paged_w4)" + ) + parser.add_argument("--var-ratio", type=float, default=0.0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--iters", type=int, default=100, help="kernel launches in the timed loop" + ) + parser.add_argument( + "--num-iters", type=int, default=101, help="run_perftest iters for --time" + ) + parser.add_argument( + "--time", action="store_true", help="also report device self-time (run_perftest)" + ) + parser.add_argument( + "--no-check", action="store_true", help="skip the one-shot correctness gate" + ) + args = parser.parse_args() + + if get_gfx() not in ("gfx942", "gfx950"): + print(f"unsupported gfx {get_gfx()}; skipping") + return + if args.preshuffle and args.kv_block_size % 16 != 0: + raise SystemExit( + f"--preshuffle requires --kv-block-size divisible by 16; " + f"got {args.kv_block_size}" + ) + + ( + q, + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + ) = build_inputs( + args.batch, + args.next_n, + args.heads, + args.head_dim, + args.kv_len, + args.kv_block_size, + DTYPE_MAP[args.q_dtype], + args.var_ratio, + args.seed, + ) + + # Oracle reads the unshuffled cache; the kernel gets the preshuffled copy. + kv_kernel = ( + preshuffle_kv_data(kv_cache_fp8, args.head_dim) + if args.preshuffle + else kv_cache_fp8 + ) + + out = torch.full( + (args.batch * args.next_n, max_model_len), + float("-inf"), + device="cuda", + dtype=torch.float32, + ) + split_kv = None if args.split_kv == 0 else args.split_kv + + def run(): + flydsl_fp8_paged_mqa_logits( + q_fp8, + kv_kernel, + weights, + out, + context_lens, + block_tables, + max_model_len, + Preshuffle=args.preshuffle, + KVBlockSize=args.kv_block_size, + ChunkK=args.chunk_k, + SplitKV=split_kv, + WavePerEU=args.wave_per_eu, + variant=args.variant, + ) + + # Warm / JIT-compile once (excluded from the timed loop and from rocprofv3 + # runs that filter on the steady-state kernel name). + run() + torch.cuda.synchronize() + + if not args.no_check: + ref = ref_fp8_paged_mqa_logits( + q, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + block_size=args.kv_block_size, + ) + neg_inf = float("-inf") + ref_mask = ref == neg_inf + diff = float( + calc_diff(out.masked_fill(out == neg_inf, 0), ref.masked_fill(ref_mask, 0)) + ) + mask_ok = bool(torch.equal(out == neg_inf, ref_mask)) + ok = (diff < 1e-3) and mask_ok + print( + f"# correctness: calc_diff={diff:.3e} mask_ok={mask_ok} " + f"-> {'PASS' if ok else 'FAIL'}" + ) + if not ok: + raise SystemExit("correctness gate failed") + + # Work estimate: QK dot is 2*H*D MACs per (query row, valid KV position). + total_ctx = int(context_lens.sum().item()) + flops = 2 * args.heads * args.head_dim * args.next_n * total_ctx + kv_bytes = args.next_n * total_ctx * (args.head_dim + 4) + + print( + f"# gfx={get_gfx()} B={args.batch} nn={args.next_n} H={args.heads} " + f"D={args.head_dim} kv_len={args.kv_len} kvb={args.kv_block_size} " + f"preshuffle={args.preshuffle} chunk_k={args.chunk_k} " + f"split_kv={args.split_kv or 'auto'} variant={args.variant or 'paged_w4'} " + f"q_dtype={args.q_dtype} iters={args.iters}" + ) + + # Timed loop: back-to-back kernel launches -- this is what rocprofv3 traces. + torch.cuda.synchronize() + for _ in range(args.iters): + run() + torch.cuda.synchronize() + + if args.time: + _, us = run_perftest(run, num_iters=args.num_iters) + us = float(us) + print( + f"time: {us:.2f} us | {flops / us / 1e6:.1f} TFLOP/s | " + f"{kv_bytes / us / 1e3:.1f} GB/s (KV read)" + ) + + +if __name__ == "__main__": + main() diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py index 08f6a322554..8d63b026665 100644 --- a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py @@ -14,6 +14,12 @@ host and the FlyDSL host auto-compute SplitKV to fill the device on small decode grids, so the comparison is apples-to-apples. +This A/B is fixed at ``KVBlockSize=1``: the Gluon reference faults on a +non-padded block-flat cache at ``KVBlockSize>1``, so any KVBlockSize>1 (incl. +Preshuffle) comparison would not be apples-to-apples. FlyDSL KVBlockSize>1 and +Preshuffle correctness is covered against a torch oracle in +``op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py`` instead. + Reuses the paged benchmark's fp8 cache packer to guarantee the same input build. """ From 92d162f744a301a4e36e3336f0e018953d4b4941 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Thu, 13 Aug 2026 14:13:57 +0000 Subject: [PATCH 09/52] refactor(flydsl): move paged mqa-logits kernel into mqa_logits subpackage Upstream relocated the mqa-logits kernels into aiter/ops/flydsl/kernels/mqa_logits/. Move the paged FP8 kernel and its shared helper module there too for consistency: kernels/fp8_paged_mqa_logits.py -> kernels/mqa_logits/fp8_paged_mqa_logits.py kernels/_mqa_logits_common.py -> kernels/mqa_logits/_mqa_logits_common.py Fix the sibling .tensor_shim imports to ..tensor_shim (now one level up) and update the public __init__ import path. No behavior change; public API (flydsl_fp8_paged_mqa_logits) is unchanged. --- aiter/ops/flydsl/__init__.py | 2 +- aiter/ops/flydsl/kernels/{ => mqa_logits}/_mqa_logits_common.py | 2 +- .../ops/flydsl/kernels/{ => mqa_logits}/fp8_paged_mqa_logits.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) rename aiter/ops/flydsl/kernels/{ => mqa_logits}/_mqa_logits_common.py (99%) rename aiter/ops/flydsl/kernels/{ => mqa_logits}/fp8_paged_mqa_logits.py (99%) diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 9aa7fb6cb14..5953c93cc3b 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -56,7 +56,7 @@ flydsl_pa_mqa_logits_fp4_prefill, flydsl_pa_mqa_logits_fp4_varqlen, ) - from .kernels.fp8_paged_mqa_logits import ( + from .kernels.mqa_logits.fp8_paged_mqa_logits import ( flydsl_fp8_paged_mqa_logits, KERNEL_VARIANTS as FP8_PAGED_MQA_LOGITS_VARIANTS, DEFAULT_VARIANT as FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT, diff --git a/aiter/ops/flydsl/kernels/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py similarity index 99% rename from aiter/ops/flydsl/kernels/_mqa_logits_common.py rename to aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index 5cfe447f452..e25fe93e085 100644 --- a/aiter/ops/flydsl/kernels/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -28,7 +28,7 @@ from flydsl.expr.numeric import ArithValue from flydsl.expr.typing import T -from .tensor_shim import GTensor, _to_raw +from ..tensor_shim import GTensor, _to_raw Vec = fx.Vector diff --git a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py similarity index 99% rename from aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py rename to aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 5e9e33d0b35..9ddd16020c9 100644 --- a/aiter/ops/flydsl/kernels/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -52,7 +52,7 @@ from flydsl._mlir.dialects import scf from flydsl._mlir import ir -from .tensor_shim import GTensor, _run_compiled, _to_raw +from ..tensor_shim import GTensor, _run_compiled, _to_raw from ._mqa_logits_common import ( MFMA_M, MFMA_N, From c3a943e5ad758375e78e4c17665c5c1030cf533b Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 14:56:48 +0000 Subject: [PATCH 10/52] fix(flydsl): size paged mqa-logits bench KV pool from real block demand Both paged benchmarks sized the block pool from max_model_len and then handed blocks out with `pool[counter % num_blocks]`, so the block table wrapped and every sequence reused the same physical blocks. At B=16, kv_len=32768 that is a 1024-block pool serving 8192 block references: an 8.65 MB KV working set standing in for 69 MB, small enough to live in cache. Every number measured against it was therefore a cache benchmark. That flattered the kernel by ~13% on its own, and distorted the head-to-head against production Gluon unevenly -- given the honest footprint FlyDSL loses 31-51% while the reference loses only 1-4%, since the reference is slow enough that memory is not its constraint. The reported speedup falls from 1.3-4.2x to 1.7-3.1x. Size the pool from the blocks the batch actually needs, build the block table vectorized rather than one element per GPU write, and report the touched footprint alongside the requested bytes. --pool-blocks shrinks the pool deliberately for cache-residency studies, and the vs-gluon bench now fails with a clear message when a shape would exceed the kernel's 2^31 i32 gather-offset limit instead of tripping an assert inside the launch path. Co-authored-by: Cursor --- .../triton/bench_flydsl_paged_mqa_logits.py | 82 ++++++++++--- .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 108 ++++++++++++++---- 2 files changed, 155 insertions(+), 35 deletions(-) diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py index 3593d2157fb..d6d23dcc08f 100644 --- a/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py @@ -34,6 +34,7 @@ from aiter.jit.utils.chip_info import get_gfx from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits +from aiter.ops.flydsl.kernels.mqa_logits.fp8_paged_mqa_logits import DEFAULT_VARIANT # Reuse the correctness test's layout/oracle helpers so the profiled kernel sees # exactly the tested co-pack + preshuffle byte layout. @@ -50,7 +51,16 @@ def build_inputs( - batch, next_n, heads, head_dim, kv_len, block_size, q_dtype, var_ratio, seed + batch, + next_n, + heads, + head_dim, + kv_len, + block_size, + q_dtype, + var_ratio, + seed, + pool_blocks=0, ): """Paged decode inputs (block-flat co-packed cache), block_size >= 1. @@ -58,13 +68,23 @@ def build_inputs( ``kv_len``: 0.0 == exact (every sequence == kv_len); v>0 draws lengths in [(1-v), (1+v)]*kv_len. A fixed ``seed`` makes the whole build deterministic (stable kernel work across profiling runs). + + The pool holds one distinct physical block per block the batch actually + needs, so the KV footprint the kernel touches is ``sum(ctx)*index_dim`` -- + what a real decode reads. Sizing the pool from ``max_model_len`` instead + (as an earlier version did) makes the block table wrap: at batch=16, + kv_len=32768, KVBlockSize=64 that is a 1024-block pool serving 8192 block + references, so all 16 sequences share the same blocks and the working set + collapses from 69 MB to 8.65 MB -- small enough to sit in cache and make + the kernel look ~13% faster than it is. ``pool_blocks`` overrides the size + to study cache residency deliberately (values below the demand re-introduce + wrapping; the padded tail of each row stays block 0 either way). """ torch.manual_seed(seed) random.seed(seed) fp8_dtype = get_fp8_e4m3_dtype() max_model_len = 2 * kv_len - num_blocks = (max_model_len + block_size - 1) // block_size if var_ratio == 0.0: context_lens = torch.full((batch,), kv_len, device="cuda", dtype=torch.int32) @@ -75,23 +95,29 @@ def build_inputs( # decode with MTP needs at least next_n tokens of context. context_lens = torch.clamp(context_lens, min=next_n) + # block table: ceil(ctx / block_size) physical blocks per sequence. + blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size + max_block_len = int(blocks_per_seq.max().item()) + needed_blocks = int(blocks_per_seq.sum().item()) + num_blocks = needed_blocks if pool_blocks <= 0 else max(pool_blocks, max_block_len) + q = torch.randn((batch, next_n, heads, head_dim), dtype=torch.bfloat16) kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) weights = torch.randn((batch * next_n, heads), dtype=torch.float32) - # block table: ceil(ctx / block_size) physical blocks per sequence. - max_block_len = (int(context_lens.max().item()) + block_size - 1) // block_size - block_tables = torch.zeros( - (batch, max_block_len), device="cuda", dtype=torch.int32 - ) + # Hand blocks out of a shuffled pool in sequence order (so a sequence's + # blocks are scattered across the pool, as they are after real paged + # allocation), padding each row's unused tail with block 0. pool = list(range(num_blocks)) random.shuffle(pool) - counter = 0 - for i in range(batch): - n_blk = (int(context_lens[i].item()) + block_size - 1) // block_size - for j in range(n_blk): - block_tables[i][j] = pool[counter % num_blocks] - counter += 1 + pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) + starts = torch.cumsum(blocks_per_seq, 0) - blocks_per_seq + col = torch.arange(max_block_len, device="cuda", dtype=torch.int64) + in_row = col[None, :] < blocks_per_seq[:, None] + draw = (starts[:, None] + col[None, :]) % num_blocks + block_tables = torch.where( + in_row, pool_t[draw], torch.zeros((), device="cuda", dtype=torch.int32) + ).to(torch.int32) q_fp8 = q.to(q_dtype) kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, fp8_dtype) @@ -139,6 +165,14 @@ def main(): ) parser.add_argument("--var-ratio", type=float, default=0.0) parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--pool-blocks", + type=int, + default=0, + help="physical blocks in the cache pool; 0 == one per block the batch " + "needs (no sharing). Smaller values make sequences share blocks, " + "shrinking the footprint into cache", + ) parser.add_argument( "--iters", type=int, default=100, help="kernel launches in the timed loop" ) @@ -181,6 +215,7 @@ def main(): DTYPE_MAP[args.q_dtype], args.var_ratio, args.seed, + args.pool_blocks, ) # Oracle reads the unshuffled cache; the kernel gets the preshuffled copy. @@ -248,15 +283,31 @@ def run(): # Work estimate: QK dot is 2*H*D MACs per (query row, valid KV position). total_ctx = int(context_lens.sum().item()) flops = 2 * args.heads * args.head_dim * args.next_n * total_ctx + # Bytes the kernel *requests*: every query row walks its whole context, so + # the next_n rows of a sequence each re-read the same blocks. kv_bytes = args.next_n * total_ctx * (args.head_dim + 4) + # Bytes the kernel *touches*: distinct blocks the table points at. Equal to + # kv_bytes/next_n unless the pool is small enough that sequences share. + n_blk_row = (context_lens.to(torch.int64) + args.kv_block_size - 1) // ( + args.kv_block_size + ) + in_row = ( + torch.arange(block_tables.shape[1], device="cuda")[None, :] < n_blk_row[:, None] + ) + distinct_blocks = int(torch.unique(block_tables[in_row]).numel()) + kv_unique = distinct_blocks * args.kv_block_size * (args.head_dim + 4) print( f"# gfx={get_gfx()} B={args.batch} nn={args.next_n} H={args.heads} " f"D={args.head_dim} kv_len={args.kv_len} kvb={args.kv_block_size} " f"preshuffle={args.preshuffle} chunk_k={args.chunk_k} " - f"split_kv={args.split_kv or 'auto'} variant={args.variant or 'paged_w4'} " + f"split_kv={args.split_kv or 'auto'} variant={args.variant or DEFAULT_VARIANT} " f"q_dtype={args.q_dtype} iters={args.iters}" ) + print( + f"# KV footprint: {kv_unique / 1e6:.1f} MB touched " + f"({distinct_blocks} blocks), {kv_bytes / 1e6:.1f} MB requested" + ) # Timed loop: back-to-back kernel launches -- this is what rocprofv3 traces. torch.cuda.synchronize() @@ -269,7 +320,8 @@ def run(): us = float(us) print( f"time: {us:.2f} us | {flops / us / 1e6:.1f} TFLOP/s | " - f"{kv_bytes / us / 1e3:.1f} GB/s (KV read)" + f"{kv_bytes / us / 1e3:.1f} GB/s requested | " + f"{kv_unique / us / 1e3:.1f} GB/s touched" ) diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py index 8d63b026665..3a5dd9959be 100644 --- a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py @@ -38,7 +38,6 @@ # same fp8 co-pack builder the production paged benchmark uses from op_tests.op_benchmarks.triton.bench_deepgemm_attention import ( kv_cache_cast_to_fp8, - cdiv, ) torch.set_default_device("cuda") @@ -103,13 +102,31 @@ def parity_mask(context_lens, batch_size, next_n, max_model_len): def build_inputs( - batch_size, next_n, heads, index_dim, kv_length, seed=0, var_ratio=0.0 + batch_size, + next_n, + heads, + index_dim, + kv_length, + seed=0, + var_ratio=0.0, + pool_blocks=0, ): """Paged input builder (blocksize==1). ``var_ratio`` controls the per-sequence context-length spread around ``kv_length``: 0.0 gives an exact length (every sequence == kv_length), and v>0 draws lengths uniformly from [(1-v)*kv_length, (1+v)*kv_length]. + + The pool holds one distinct block per block the batch needs, so the KV + footprint both kernels read is ``sum(ctx)*(index_dim+4)`` -- what a real + decode touches. Sizing it from ``max_model_len`` instead (as an earlier + version did) makes the block table wrap: at blocksize=1 the pool is + 2*kv_length blocks while the batch needs batch_size*kv_length, so from + batch_size=4 up the sequences share blocks and the working set collapses + into cache. At B=16, kv_length=32768 that is 8.65 MB standing in for 69 MB, + which flatters both kernels and, because the two spend their time + differently, does not flatter them equally. ``pool_blocks`` overrides the + size to study cache residency deliberately. """ torch.manual_seed(seed) random.seed(seed) @@ -117,7 +134,6 @@ def build_inputs( max_model_len = 2 * kv_length blocksize = 1 - num_blocks = (max_model_len + blocksize - 1) // blocksize if var_ratio == 0.0: context_lens = torch.full( @@ -136,6 +152,21 @@ def build_inputs( # decode with MTP needs at least next_n tokens of context. context_lens = torch.clamp(context_lens, min=next_n) + blocks_per_seq = (context_lens.to(torch.int64) + blocksize - 1) // blocksize + max_block_len = int(blocks_per_seq.max().item()) + needed_blocks = int(blocks_per_seq.sum().item()) + num_blocks = needed_blocks if pool_blocks <= 0 else max(pool_blocks, max_block_len) + + # The kernels ride the per-token byte offset on a 32-bit buffer voffset, so + # the whole pool must be addressable in i32 bytes. + pool_bytes = num_blocks * blocksize * (index_dim + 4) + if pool_bytes >= 2**31: + raise SystemExit( + f"pool of {num_blocks} blocks = {pool_bytes / 1e9:.2f} GB exceeds the " + f"i32 gather-offset limit (2.15 GB). Shrink the shape or cap the pool " + f"with --pool-blocks (which re-introduces block sharing)." + ) + q = torch.randn( (batch_size, next_n, heads, index_dim), device="cuda", dtype=torch.bfloat16 ) @@ -146,20 +177,18 @@ def build_inputs( (batch_size * next_n, heads), device="cuda", dtype=torch.float32 ) - max_block_len = (context_lens.max().item() + blocksize - 1) // blocksize * blocksize - block_tables = torch.zeros( - (batch_size, max_block_len), device="cuda", dtype=torch.int32 - ) + # Hand blocks out of a shuffled pool in sequence order (scattered across the + # pool, as they are after real paged allocation); pad each row's tail with 0. pool = list(range(num_blocks)) random.shuffle(pool) pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) - counter = 0 - for i in range(batch_size): - ctx_len = int(context_lens[i].item()) - n = cdiv(ctx_len, blocksize) - idx = (counter + torch.arange(n, device="cuda")) % num_blocks - block_tables[i, :n] = pool_t[idx] - counter += n + starts = torch.cumsum(blocks_per_seq, 0) - blocks_per_seq + col = torch.arange(max_block_len, device="cuda", dtype=torch.int64) + in_row = col[None, :] < blocks_per_seq[:, None] + draw = (starts[:, None] + col[None, :]) % num_blocks + block_tables = torch.where( + in_row, pool_t[draw], torch.zeros((), device="cuda", dtype=torch.int32) + ).to(torch.int32) q_fp8 = q.to(fp8_dtype) kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, padding=False, fp8_dtype=fp8_dtype) @@ -172,6 +201,7 @@ def build_inputs( block_tables, max_model_len, fp8_dtype, + min(needed_blocks, num_blocks) * blocksize * (index_dim + 4), ) @@ -184,7 +214,15 @@ def time_kernel(fn, repeats, num_iters): def run_shape( - batch_size, next_n, heads, index_dim, kv_length, repeats, num_iters, var_ratio=0.0 + batch_size, + next_n, + heads, + index_dim, + kv_length, + repeats, + num_iters, + var_ratio=0.0, + pool_blocks=0, ): ( q, @@ -195,8 +233,15 @@ def run_shape( block_tables, max_model_len, fp8_dtype, + touched_bytes, ) = build_inputs( - batch_size, next_n, heads, index_dim, kv_length, var_ratio=var_ratio + batch_size, + next_n, + heads, + index_dim, + kv_length, + var_ratio=var_ratio, + pool_blocks=pool_blocks, ) ref = ref_fp8_paged_mqa_logits( @@ -275,6 +320,7 @@ def fn_fly(): "H": heads, "D": index_dim, "avg_kv": kv_length, + "kv_mb": touched_bytes / 1e6, "ref_ms": ref_med / 1e3, "ref_min_ms": ref_min / 1e3, "ref_max_ms": ref_max / 1e3, @@ -318,6 +364,14 @@ def main(): help="Per-sequence context-length spread around kv_len. 0.0 = exact " "(every sequence == kv_len); e.g. 0.5 draws lengths in [0.5, 1.5]*kv_len.", ) + parser.add_argument( + "--pool-blocks", + type=int, + default=0, + help="physical blocks in the cache pool; 0 == one per block the batch " + "needs (no sharing). Smaller values make sequences share blocks, " + "shrinking the KV footprint into cache", + ) args = parser.parse_args() if get_gfx() not in ("gfx942", "gfx950"): @@ -329,21 +383,35 @@ def main(): f"# arch={get_gfx()} fp8={get_fp8_e4m3_dtype()} " f"repeats={args.repeats} num_iters={args.num_iters} " f"ref_chunk_k={REF_CHUNK_K} ref_wave_per_eu={REF_WAVE_PER_EU} " - f"ctx_len={len_mode} max_model_len=2*kv_len blocksize=1" + f"ctx_len={len_mode} max_model_len=2*kv_len blocksize=1 " + f"pool={'one block per block needed' if args.pool_blocks <= 0 else args.pool_blocks}" + ) + print( + "# kv_mb = distinct KV bytes the batch touches " + "(both kernels read it once per query row)" ) header = ( - "B nn grid H D kv_len | ref_ms[min-max] fly_ms[min-max] fly/ref | " + "B nn grid H D kv_len kv_mb | ref_ms[min-max] fly_ms[min-max] fly/ref | " "ref_diff fly_diff ref_mask fly_mask ref_pass fly_pass" ) print(header) rows = [] for B, nn, H, D, kv_len in default_shapes(): r = run_shape( - B, nn, H, D, kv_len, args.repeats, args.num_iters, var_ratio=args.var_ratio + B, + nn, + H, + D, + kv_len, + args.repeats, + args.num_iters, + var_ratio=args.var_ratio, + pool_blocks=args.pool_blocks, ) rows.append(r) print( - f"{r['B']} {r['nn']} {r['grid']} {r['H']} {r['D']} {r['avg_kv']} | " + f"{r['B']} {r['nn']} {r['grid']} {r['H']} {r['D']} {r['avg_kv']} " + f"{r['kv_mb']:.0f} | " f"{r['ref_ms']:.4f}[{r['ref_min_ms']:.4f}-{r['ref_max_ms']:.4f}] " f"{r['fly_ms']:.4f}[{r['fly_min_ms']:.4f}-{r['fly_max_ms']:.4f}] " f"{r['fly/ref']:.2f}x | " From e28ac81fec2f1f78ab20011aaa082d1b2fe13b76 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 14:57:04 +0000 Subject: [PATCH 11/52] perf(flydsl): retune paged mqa-logits launch geometry (9.3% across shapes) SplitKV was auto-computed with the production Triton host formula, which ignores how many waves a CTA contains and so drifts with shape: it lands 2.5 waves per SIMD at B=16 but 10 at B=128. Resident wave count is the lever that matters here -- this kernel is bound by memory throughput rather than per-wave latency, so too few waves leaves the memory system idle and too many just queue. Measured across six shapes (batch 4..128, kv_len 4k..128k) the curve is flat between 2.0 and 2.5 waves/SIMD and rises steeply either side: 1.0/SIMD costs 33.9us and 4.0/SIMD 29.4us against 27.6us at the optimum. Target WavePerEU waves per SIMD directly, accounting for the CTA's wave count, and default to paged_w1. At equal resident wave count, splitting a CTA across 2 or 4 waves is slower on every shape measured: a 1-wave CTA schedules more freely, and the waves have no work to share anyway since each owns a disjoint set of KV columns with no barrier or LDS between them. Worth 4.6-14.4% per shape, 9.3% aggregate, and it removes the shape sensitivity -- all six shapes now land within 1% of 27.7us for the same total KV work, where the old defaults spread them over 29.2-32.5us. Verified against the full 870-case correctness sweep (max error 5.3e-4, tolerance 1e-3); the block-flat KVBlockSize=1 path is unchanged. Two public-surface notes for reviewers. This deliberately diverges from the deepgemm_fp8_paged_mqa_logits host formula it previously mirrored. And WavePerEU changes meaning, from a multiplier in the old formula to the target waves per SIMD; its default of 2 is where the curve bottoms out. FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT changes paged_w4 -> paged_w1. Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits.py | 56 ++++++++++++++----- 1 file changed, 42 insertions(+), 14 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 9ddd16020c9..49cadfc86b9 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -67,18 +67,33 @@ ) -def _auto_split_kv(batch_size, next_n, wave_per_eu, device_index, total_cu=None): - """Production host SplitKV formula (mirrors deepgemm_fp8_paged_mqa_logits): +_SIMDS_PER_CU = 4 - SplitKV = ((max(1, TotalCuCount // (batch*next_n)) + 4) // 5 * 5) * WavePerEU - Fills the device for small decode grids where the batch*next_n row grid alone - leaves most CUs idle. Returns >= 1. +def _auto_split_kv( + batch_size, next_n, wave_per_eu, device_index, total_cu=None, waves_per_block=1 +): + """Pick SplitKV so the dispatch lands ``wave_per_eu`` waves on every SIMD: + + SplitKV = round(WavePerEU * 4 * TotalCuCount / (batch*next_n*WavePerBlock)) + + Resident wave count is the lever that matters here -- the kernel is bound by + memory throughput, not per-wave latency, so too few waves leaves the memory + system idle and too many just queue. Measured on gfx950/256 CU across six + shapes (batch 4..128, kv_len 4k..128k), the curve is flat between 2.0 and 2.5 + waves/SIMD and rises steeply either side: at the reference shape 1.0/SIMD + costs 33.9 us and 4.0/SIMD costs 29.4 us against 27.6 us at 2.0. + + The earlier formula, ``((TotalCuCount // (batch*next_n) + 4) // 5 * 5) * + WavePerEU``, ignored the CTA's wave count, so it drifted with shape -- + 2.5 waves/SIMD at the reference shape but 10 at batch=128 -- and left 4-14% + on the table. Returns >= 1. """ tile_q_count = max(1, batch_size * next_n) if total_cu is None: total_cu = device_cu_count(device_index) - split_kv = ((max(1, total_cu // tile_q_count) + 4) // 5 * 5) * wave_per_eu + target_waves = max(1, wave_per_eu) * _SIMDS_PER_CU * total_cu + split_kv = round(target_waves / (tile_q_count * max(1, waves_per_block))) return max(1, int(split_kv)) @@ -401,8 +416,14 @@ def launch_fp8_paged_mqa_logits( # Kernel variants: single-token-per-block, wave-split only ("paged_w"). # WPB must divide the column-tile count BKV/16 (=8 at the default BKV=128). +# +# WPB=1 (one wave per CTA) is the default: with SplitKV set to hit the same +# resident wave count, splitting a CTA across 2 or 4 waves is 4-14% slower on +# every shape measured -- a 1-wave CTA schedules more freely and pays no +# whole-CTA granularity cost, and the waves have no work to share anyway (each +# owns a disjoint set of KV columns, with no barrier or LDS between them). KERNEL_VARIANTS = tuple(f"paged_w{w}" for w in (1, 2, 4)) -DEFAULT_VARIANT = "paged_w4" +DEFAULT_VARIANT = "paged_w1" def _variant_wpb(variant): @@ -503,11 +524,13 @@ def flydsl_fp8_paged_mqa_logits( per-token page table). max_model_len: int, out_logits column count. SplitKV: KV-column split count (grid = split_kv*batch*next_n). None => - the production host formula (fills the device on small decode - grids); pass an int to override, 1 disables splitting. Splits - own disjoint, gap-free BKV-aligned column ranges (one writer - per column, no cross-CTA reduction). - WavePerEU/TotalCuCount: inputs to the auto-SplitKV formula. + sized so the dispatch lands WavePerEU waves on every SIMD; + pass an int to override, 1 disables splitting. Splits own + disjoint, gap-free BKV-aligned column ranges (one writer per + column, no cross-CTA reduction). + WavePerEU: target resident waves per SIMD for the auto-SplitKV formula + (default 2, which is where the measured curve bottoms out). + TotalCuCount: CU count for that formula; None => query the device. ChunkK: KV tile width (must be a multiple of MFMA_N=32). Mirrors the Triton ChunkK parameter. Defaults to 128. @@ -591,10 +614,15 @@ def flydsl_fp8_paged_mqa_logits( kv_bytes = kv_cache.reshape(-1) # KV-column splits (one lever): fill the device when the batch*next_n row - # grid is small. Auto == production host formula; overridable. + # grid is small. Auto targets WavePerEU waves per SIMD; overridable. if SplitKV is None: split_kv = _auto_split_kv( - batch_size, next_n, WavePerEU, q_fp8.device.index, TotalCuCount + batch_size, + next_n, + WavePerEU, + q_fp8.device.index, + TotalCuCount, + waves_per_block=_variant_wpb(variant), ) else: split_kv = max(1, int(SplitKV)) From 2341f1ffdb24440cdd9fefd7e9cbd856d40dcbcf Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 15:19:39 +0000 Subject: [PATCH 12/52] refactor(flydsl): use FlyDSL operator sugar in paged mqa-logits The kernel was written against raw MLIR builders -- 58 _to_raw() calls and 65 arith.* calls -- for arithmetic FlyDSL already overloads. The worst case was the store predicate, 17 lines of nested arith.andi/arith.cmpi to say `(lane_div_N == 0) & (col <= q_limit)`. Route unsigned division and modulo through new udiv/umod helpers rather than the bare `//` and `%` operators: fx.Int32 is signed, so `//` lowers to floordivsi, which expands to several instructions to get negative-operand rounding right. Every quantity divided here is non-negative, so going via fx.Uint32 keeps the single-instruction divui/remui the hand-written calls emitted. Verified codegen-neutral: the ATT-decoded ISA is byte-identical to the previous commit, 655 instructions in the same order. 870/870 correctness cases pass at 5.3e-4, timing is unchanged at 28.0us, and ruff reports the same findings as before. Co-authored-by: Cursor --- .../kernels/mqa_logits/_mqa_logits_common.py | 25 +++++++--- .../mqa_logits/fp8_paged_mqa_logits.py | 50 ++++++------------- 2 files changed, 32 insertions(+), 43 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index e25fe93e085..b32c4a4e1f4 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -73,6 +73,19 @@ def i32_add(a, b): return fx.Int32(arith.addi(_to_raw(a), _to_raw(b))) +# `//` and `%` on a signed Int32 lower to floordivsi/remsi, which expand to +# several instructions to get negative-operand rounding right. Every quantity +# divided here (lane ids, block ids, byte offsets) is non-negative, so route +# through Uint32 to keep the single-instruction divui/remui the hand-written +# arith calls used to emit. +def udiv(a, b): + return fx.Int32(fx.Uint32(a) // fx.Uint32(b)) + + +def umod(a, b): + return fx.Int32(fx.Uint32(a) % fx.Uint32(b)) + + @lru_cache(maxsize=8) def device_cu_count(device_index: int) -> int: """Compute-unit count for a CUDA/HIP device (cached); 304 if unavailable.""" @@ -109,10 +122,7 @@ def load_pack_v8i32(i32_view, byte_off_i32, lane8): """ def _load_i64(off): - dw = fx.Int32( - arith.divui(_to_raw(byte_off_i32 + lane8 + off), _to_raw(fx.Int32(4))) - ) - v2 = i32_view.vec_load((dw,), vec_size=2) + v2 = i32_view.vec_load((udiv(byte_off_i32 + lane8 + off, 4),), vec_size=2) return Vec(v2).bitcast(fx.Int64)[0].ir_value() i64_0 = _load_i64(0) # K-group 0: k = lane_div_N*8 + 0..7 @@ -151,13 +161,12 @@ def load_pack_v8i32_preshuffle( D = head_size # Token super-row (16-token blocks) + intra-super-row column; loop-invariant # across the four hidden groups, so hoist out of the per-group load. - c_hi = fx.Int32(arith.divui(_to_raw(tok_in_block), _to_raw(fx.Int32(16)))) - c_lo = fx.Int32(arith.remui(_to_raw(tok_in_block), _to_raw(fx.Int32(16)))) + c_hi = udiv(tok_in_block, 16) + c_lo = umod(tok_in_block, 16) base = block_byte_base + c_hi * (D * 16) + c_lo * 16 + lane8 def _load_i64(g): - dw = fx.Int32(arith.divui(_to_raw(base + g * 256), _to_raw(fx.Int32(4)))) - v2 = i32_view.vec_load((dw,), vec_size=2) + v2 = i32_view.vec_load((udiv(base + g * 256, 4),), vec_size=2) return Vec(v2).bitcast(fx.Int64)[0].ir_value() g0 = kk_step * 4 # first hidden-16 group of this MFMA K-step diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 49cadfc86b9..05a3d87f52b 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -64,6 +64,8 @@ load_pack_v8i32_preshuffle, make_out_row_view, mfma_head_reduce, + udiv, + umod, ) @@ -196,15 +198,15 @@ def kernel( # ChunkQ==H => grid = split_kv * batch * next_n (split outermost), # mirroring the Triton _deepgemm_fp8_paged_mqa_logits pid decomposition. - pid_next_n = fx.Int32(arith.remui(_to_raw(bid), _to_raw(next_n))) - _rem = fx.Int32(arith.divui(_to_raw(bid), _to_raw(next_n))) - pid_batch = fx.Int32(arith.remui(_to_raw(_rem), _to_raw(batch_size))) - pid_split_kv = fx.Int32(arith.divui(_to_raw(_rem), _to_raw(batch_size))) - - wave = fx.Int32(arith.divui(_to_raw(tid), _to_raw(fx.Int32(64)))) - lane = fx.Int32(arith.remui(_to_raw(tid), _to_raw(fx.Int32(64)))) - lane_div_N = fx.Int32(arith.divui(_to_raw(lane), _to_raw(fx.Int32(MFMA_N)))) - lane_mod_N = fx.Int32(arith.remui(_to_raw(lane), _to_raw(fx.Int32(MFMA_N)))) + pid_next_n = umod(bid, next_n) + _rem = udiv(bid, next_n) + pid_batch = umod(_rem, batch_size) + pid_split_kv = udiv(_rem, batch_size) + + wave = udiv(tid, 64) + lane = umod(tid, 64) + lane_div_N = udiv(lane, MFMA_N) + lane_mod_N = umod(lane, MFMA_N) lane8 = lane_div_N * 8 # fp8 operands loaded as vector<8xi32> per K-step (32 bytes = 4 i64s) @@ -294,12 +296,8 @@ def kernel( # away (block_idx==col_c, tok_in_block==0) -- no Python `if` here # because the FlyDSL AST rewriter would turn it into a runtime # scf.if and drop the branch-local addresses. - block_idx = fx.Int32( - arith.divui(_to_raw(col_c), _to_raw(fx.Int32(KVB))) - ) - tok_in_block = fx.Int32( - arith.remui(_to_raw(col_c), _to_raw(fx.Int32(KVB))) - ) + block_idx = udiv(col_c, KVB) + tok_in_block = umod(col_c, KVB) ind_off = pid_batch * max_block_len + block_idx physical = fx.Int32(ind_t[ind_off]) block_byte_base = physical * block_byte_stride @@ -324,10 +322,7 @@ def kernel( # Co-packed per-token f32 scale (block-flat: after the KVB keys). scale_byte = block_byte_base + (KVB * D) + tok_in_block * 4 - scale_dword = fx.Int32( - arith.divui(_to_raw(scale_byte), _to_raw(fx.Int32(4))) - ) - kv_scale = _to_raw(fx.Float32(kv_f32[scale_dword])) + kv_scale = _to_raw(fx.Float32(kv_f32[udiv(scale_byte, 4)])) if scale_mul != 1.0: kv_scale = arith.MulFOp( kv_scale, scale_mul_c, fastmath=fm_fast @@ -349,22 +344,7 @@ def kernel( # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. # Write only in-causal columns; masked positions keep the # caller's -inf prefill (matches clean_logits semantics). - is_writer = arith.andi( - _to_raw( - arith.cmpi( - arith.CmpIPredicate.eq, - _to_raw(lane_div_N), - _to_raw(fx.Int32(0)), - ) - ), - _to_raw( - arith.cmpi( - arith.CmpIPredicate.sle, - _to_raw(col), - _to_raw(q_limit), - ) - ), - ) + is_writer = _to_raw((lane_div_N == 0) & (col <= q_limit)) with ir.InsertionPoint(scf.IfOp(is_writer).then_block): out_row_t[col] = fx.Float32(col_sum) scf.YieldOp([]) From e415028bb8a1ed9a178e38c2dfc2f5758683262f Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 15:26:38 +0000 Subject: [PATCH 13/52] refactor(flydsl): drop dead helpers and tighten mqa-logits docs _mqa_logits_common carried three helpers nothing referenced: fn_to_fnuz_i64 (45 lines -- the dense kernel defines its own local copy), load_pack_i64 (likewise), and i32_add. Removing them is the bulk of the shrink. Comments and docstrings were 30% of the two files, much of it restating the code or duplicating the PR description. Trimmed to what a reader cannot derive: the co-pack layout contract, the ReLU scale-hoist identity, the i32 gather ceiling, the D-reg-to-head mapping, and the warning that a Python `if` on KVB would become a runtime scf.if and drop branch-local addresses. 932 -> 770 lines across the two files. ISA byte-identical to the previous commit (655 instructions), 870/870 correctness at 5.3e-4, 27.8us unchanged, no new ruff findings. Co-authored-by: Cursor --- .../kernels/mqa_logits/_mqa_logits_common.py | 172 ++++-------------- .../mqa_logits/fp8_paged_mqa_logits.py | 159 ++++++---------- 2 files changed, 91 insertions(+), 240 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index b32c4a4e1f4..52b693c16a6 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -1,19 +1,14 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Shared FlyDSL building blocks for the FP8 MQA-logits kernels. +"""Shared FlyDSL building blocks for the dense and paged FP8 MQA-logits kernels. -Both the dense (``fp8_mqa_logits.py``) and paged (``fp8_paged_mqa_logits.py``) -kernels compute the same core per query row / KV column:: +Both compute the same core per query row / KV column:: logits[row, n] = sum_h ReLU( * kv_scale[n]) * weights[row, h] -They differ only in how K / kv_scale are addressed (contiguous vs paged gather) -and in the windowing / output layout. Everything that is identical -- the fp8 -16x16x32 MFMA compute, the ReLU * weight head-sum + kv-scale hoist, the in-wave -``shuffle_xor`` head reduce, the fp8 dword-pack load, the FN->FNUZ byte patch, -and the i64 byte-base per-row output view -- lives here so neither kernel -duplicates it. +and differ only in how K / kv_scale are addressed (contiguous vs paged gather) +and in the windowing / output layout. Everything else lives here. """ # No `from __future__ import annotations`: FlyDSL arg typing needs real @@ -32,28 +27,21 @@ Vec = fx.Vector -# MFMA tile dims: 32x32x64 scaled fp8 (CDNA4, gfx950). -# K=64 per step → K_STEPS = head_size // 64 (2 for D=128). -# A/B operands are vector<8xi32> (4 packed i64s = 32 fp8 bytes per lane). -# 16 f32 D-regs/lane, same output tile shape as 32x32x16 (same D-reg layout). -# lane_div_N = lane // 32 (0 or 1), lane_mod_N = lane % 32 (KV col index). -# Scale operands use neutral e8m0 (0x7F7F7F7F = 1.0 per byte, 4 packed). +# 32x32x64 scaled fp8 MFMA (CDNA4). A/B are vector<8xi32> per lane (32 fp8 +# bytes), C/D are 16 f32 D-regs. lane_div_N = lane//32, lane_mod_N = lane%32. MFMA_M = 32 MFMA_N = 32 MFMA_K = 64 -# The scaled 32x32x64 MFMA: A/B are vector<8xi32>, same D-reg layout as 32x32x16. -# operands: [a128, b128, c, cbsz=0, blgp=0, opsel_a=0, scale_a, opsel_b=0, scale_b] -# where a128/b128 are Vec(8xi32) (4 i64s packed), scale_a/scale_b are i32 values -# with 4 packed e8m0 bytes. Neutral: 0x7F7F7F7F (all 1.0). def _mfma_scale_f32_32x32x64_fp8(result_type, operands): + """operands: [a, b, c, cbsz, blgp, opsel_a, scale_a, opsel_b, scale_b].""" from flydsl.expr.rocdl import _unwrap_mfma_operand - a = _unwrap_mfma_operand(operands[0]) # vector<8xi32> - b = _unwrap_mfma_operand(operands[1]) # vector<8xi32> - c = _unwrap_mfma_operand(operands[2]) # vector<16xf32> - # Neutral scale: 0x7F7F7F7F = 4 packed e8m0 bytes, each 127 → 2^(127-127)=1.0 + a = _unwrap_mfma_operand(operands[0]) + b = _unwrap_mfma_operand(operands[1]) + c = _unwrap_mfma_operand(operands[2]) + # Neutral e8m0 scale: 4 packed bytes of 127 => 2^(127-127) = 1.0 each. neutral = arith.constant(0x7F7F7F7F, type=T.i32) return rocdl.mfma_scale_f32_32x32x64_f8f6f4_( result_type, a, b, c, 0, 0, 0, neutral, 0, neutral @@ -68,16 +56,9 @@ def _mfma_scale_f32_32x32x64_fp8(result_type, operands): } -def i32_add(a, b): - """i32 add (result stays Int32, not index type).""" - return fx.Int32(arith.addi(_to_raw(a), _to_raw(b))) - - -# `//` and `%` on a signed Int32 lower to floordivsi/remsi, which expand to -# several instructions to get negative-operand rounding right. Every quantity -# divided here (lane ids, block ids, byte offsets) is non-negative, so route -# through Uint32 to keep the single-instruction divui/remui the hand-written -# arith calls used to emit. +# Via Uint32 because `//` on a signed Int32 lowers to floordivsi, which expands +# to several instructions for negative-operand rounding. Everything divided here +# is non-negative, so this keeps the single-instruction divui/remui. def udiv(a, b): return fx.Int32(fx.Uint32(a) // fx.Uint32(b)) @@ -95,30 +76,11 @@ def device_cu_count(device_index: int) -> int: return 304 -def load_pack_i64(i32_view, byte_off_i32): - """Load 8 fp8 bytes as 2 i32 dwords and bitcast to i64. - - fp8 operands are read 8 bytes at a time as 2 i32 dwords (a ``v8i8`` - buffer_load fails to lower on gfx942), then bitcast to i64 for the MFMA. - ``i32_view`` is a ``GTensor(..., dtype=T.i32)`` over the fp8 bytes and - ``byte_off_i32`` is the byte offset of the first of the 8 bytes. - """ - dword_off = fx.Int32(arith.divui(_to_raw(byte_off_i32), _to_raw(fx.Int32(4)))) - v2 = i32_view.vec_load((dword_off,), vec_size=2) - return Vec(v2).bitcast(fx.Int64)[0].ir_value() - - def load_pack_v8i32(i32_view, byte_off_i32, lane8): - """Load 32 fp8 bytes as vector<8xi32> for the 32x32x64 scaled MFMA. - - The 32x32x64 MFMA (K=64) A/B operands are vector<8xi32> per lane. - Each lane holds 32 fp8 bytes covering 4 consecutive K-groups of 16. - Same layout as 4 consecutive 32x32x16 K-steps: the per-lane K-bytes are - at positions [lane_div_N*8 + kk*16 + 0..7] for kk in 0..3. + """Load 32 fp8 bytes as the vector<8xi32> A/B operand of the 32x32x64 MFMA. - ``byte_off_i32`` is the base byte offset (tok_byte or q_row_base + h_a*stride) - for the first K-group. ``lane8 = lane_div_N * 8`` is the within-K-group offset. - Returns a vector<8xi32> value (MLIR Value). + Each lane holds 4 K-groups of 8 bytes at ``lane8 + kk*16``. Reads go out as + 2-dword loads because a ``v8i8`` buffer_load fails to lower on gfx942. """ def _load_i64(off): @@ -135,32 +97,14 @@ def _load_i64(off): def load_pack_v8i32_preshuffle( i32_view, block_byte_base, tok_in_block, head_size, kk_step, lane8 ): - """Load 32 fp8 bytes as vector<8xi32> from the production preshuffled layout. - - The production ``Preshuffle`` path stores each physical block's fp8 key data - reordered by ``shuffle_weight(layout=(16,16))`` (only the ``KVB*D`` data - section; the co-packed f32 scale tail is left token-ordered, so the scale - gather is unchanged). Element ``(token c, hidden h)`` lands at the intra-block - byte offset:: - - off(c,h) = (c//16)*(D*16) + (h//16)*256 + (c%16)*16 + (h%16) - - which is exactly the inverse of the host shuffle permute (and matches the - gluon CDNA ``offset_k_fixed``). This kernel's per-lane K payload for MFMA - K-step ``kk_step`` is the four contiguous hidden-16 groups - ``g = kk_step*4 + sub`` (``sub`` in 0..3); within a group the 8-byte half - picked by ``lane8`` (``= lane_div_N*8``) is still contiguous, so we issue the - same four ``dwordx2`` loads and assemble a ``vector<8xi32>``. - - ``block_byte_base`` is the physical block's byte base in the cache pool - (``physical * KVB * index_dim``); the shuffled data section starts at - offset 0 within it. ``head_size`` (``D``) must be a multiple of 32 (the 16x16 - shuffle pairs ``k_width=16`` fp8 into 32-byte chunks) -- guaranteed here by - ``D % MFMA_K == 0``. + """Same operand as ``load_pack_v8i32``, from the ``shuffle_weight(16,16)`` + layout the production path stores. Element (token c, hidden h) of a block + sits at ``(c//16)*(D*16) + (h//16)*256 + (c%16)*16 + (h%16)``, the inverse + of the host permute. Only the ``KVB*D`` data section is shuffled; the f32 + scale tail stays token-ordered, so the scale gather is unchanged. """ D = head_size - # Token super-row (16-token blocks) + intra-super-row column; loop-invariant - # across the four hidden groups, so hoist out of the per-group load. + # Loop-invariant across the four hidden groups, so hoisted out of the load. c_hi = udiv(tok_in_block, 16) c_lo = umod(tok_in_block, 16) base = block_byte_base + c_hi * (D * 16) + c_lo * 16 + lane8 @@ -177,56 +121,10 @@ def _load_i64(g): return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) -def fn_to_fnuz_i64(raw_i64): - """Map FN byte 0x80 (neg-zero) -> 0x00 in 8 packed fp8 bytes. - - An e4m3 FN (OCP, bias 7) operand encodes exactly 2x the value the gfx942 - FNUZ (bias 8) MFMA reads for the same byte, except FN -0 = 0x80 which is - FNUZ NaN; patch that byte to +0 and let the caller undo the 2x via the - kv_scale (ReLU positive-homogeneity keeps the result exact). - """ - lo_i32 = arith.TruncIOp(T.i32, raw_i64).result - hi_i64 = arith.ShRUIOp(raw_i64, arith.constant(32, type=T.i64)).result - hi_i32 = arith.TruncIOp(T.i32, hi_i64).result - - def _fix_i32(src): - result = arith.constant(0, type=T.i32) - for byte_idx in range_constexpr(4): - shift = arith.constant(byte_idx * 8, type=T.i32) - byte_val = arith.andi( - arith.shrui(src, shift), - arith.constant(0xFF, type=T.i32), - ) - is_0x80 = arith.cmpi( - arith.CmpIPredicate.eq, - byte_val, - arith.constant(0x80, type=T.i32), - ) - cleaned = arith.select( - is_0x80, - arith.constant(0, type=T.i32), - byte_val, - ) - result = arith.ori(result, arith.shli(cleaned, shift)) - return result - - lo_fix = _fix_i32(lo_i32) - hi_fix = _fix_i32(hi_i32) - lo_64 = arith.ExtUIOp(T.i64, lo_fix).result - hi_64 = arith.ShLIOp( - arith.ExtUIOp(T.i64, hi_fix).result, arith.constant(32, type=T.i64) - ).result - return arith.OrIOp(lo_64, hi_64).result - - def make_out_row_view(logits, stride_i64, row_i32): - """1-D output GTensor for ``row_i32``; byte base computed in i64. - - The row's i64 byte offset (``row * stride * 4`` bytes for f32) goes into the - base pointer so the remaining column offset stays in i32. A 2-D (row, col) - view would compute ``row * stride + col`` in i32 and overflow past 2^31 - (~46k-square dense outputs / large paged ``max_model_len``), silently - mis-writing. + """1-D output GTensor for ``row_i32``, with the row byte base folded into + the pointer in i64 so the column offset stays i32. A 2-D (row, col) view + would compute ``row*stride + col`` in i32 and silently overflow past 2^31. """ _ri64 = arith.extui(T.i64, _to_raw(row_i32)) _byte = arith.muli(arith.muli(_ri64, stride_i64), arith.constant(4, type=T.i64)) @@ -248,18 +146,12 @@ def mfma_head_reduce( mfma_fn=MFMA_FN, dreg_count=16, ): - """One column's logit: fp8 MFMA over heads, ReLU * weight sum, kv-scale, head reduce. - - ``a_row`` : ``a_row[mi][kk]`` Q fragments (i64 or vector<8xi32>) for this row. - ``b_col`` : ``b_col[kk]`` K fragments (same type as a_row) for this KV tile. - ``w_row`` : ``w_row[mi][ii]`` f32 head weights for this query row. - ``kv_scale``: f32 per-KV-token dequant scale (>= 0), hoisted out of the - head sum (ReLU is positive-homogeneous, so ReLU(s*x)=s*ReLU(x) - and the whole column sum is scaled once instead of per head). - ``dreg_count``: number of f32 D-regs per lane (16 for 32x32 tiles). - - Returns the reduced f32 ``col_sum`` (raw ir value). The in-wave (width 64) - head reduce uses shuffle_xor(32) for the 32x32 MFMA (lane_div_N in {0,1}). + """One column's logit: fp8 MFMA over heads, ReLU * weight sum, kv-scale, + in-wave head reduce. Returns the reduced f32 ``col_sum`` (raw ir value). + + ``kv_scale`` (>= 0) is applied once to the column sum rather than per head: + ReLU is positively homogeneous, so ReLU(s*x) = s*ReLU(x) and hoisting it out + of the head sum is exact. """ col_sum = _to_raw(f32_0) for mi in range_constexpr(m_tiles): diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 05a3d87f52b..ae6fbcea067 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -14,25 +14,19 @@ maps through the block table (``kv_indices``) to a physical block, and each token's fp8 K bytes and its f32 dequant scale are **co-packed** in that block. -Supported scope (this file): ``KVBlockSize >= 1`` (block-flat co-packed layout) -with SplitKV KV-column parallelism. The public ``flydsl_fp8_paged_mqa_logits`` -mirrors the tensor contract of the Triton -``aiter.ops.triton.attention.pa_mqa_logits.deepgemm_fp8_paged_mqa_logits`` so the -two are interchangeable in tests/benchmarks. ``Preshuffle=True`` consumes the -production ``shuffle_weight(layout=(16,16))`` KV data layout (requires -``KVBlockSize % 16 == 0``). - -Cache layout (per physical block of ``KVBlockSize`` tokens): all ``KVBlockSize`` -fp8 key rows (``D`` bytes each) grouped first, then all ``KVBlockSize`` f32 -dequant scales (block-flat). The per-lane gather resolves each logical KV column -independently through the block table (``kv_indices[b, p // KVBlockSize]``) so a -column tile may straddle block boundaries; ``KVBlockSize == 1`` degenerates to -the per-token ``[D fp8 | 4 scale]`` slot layout. - -The fp8 32x32x64 scaled MFMA compute, ReLU*weight head-sum + kv-scale hoist, -in-wave ``shuffle_xor`` head reduce, the fp8 dword-pack load, and the i64 -byte-base per-row output view are all shared with the dense kernel via -``._mqa_logits_common`` (no duplication). +Supports ``KVBlockSize >= 1`` with SplitKV KV-column parallelism, and mirrors +the tensor contract of the Triton ``deepgemm_fp8_paged_mqa_logits`` so the two +are interchangeable. ``Preshuffle=True`` consumes the production +``shuffle_weight(layout=(16,16))`` KV layout (needs ``KVBlockSize % 16 == 0``). + +Cache layout, per physical block of ``KVBlockSize`` tokens: the fp8 key rows +(``D`` bytes each) grouped first, then the f32 dequant scales. Each lane +resolves its own column through ``kv_indices[b, p // KVBlockSize]``, so a column +tile may straddle blocks; ``KVBlockSize == 1`` degenerates to per-token +``[D fp8 | 4 scale]`` slots. + +The MFMA compute, ReLU*weight head-sum, head reduce, dword-pack load and output +view are shared with the dense kernel via ``._mqa_logits_common``. """ # No `from __future__ import annotations`: FlyDSL arg typing needs real @@ -115,28 +109,17 @@ def _build_paged_kernel( convert_q_fn: bool = False, convert_kv_fn: bool = False, ): - """Paged MQA-logits kernel (KVBlockSize>=1) with SplitKV column parallelism. - - One thread block owns one ``(batch, next_n)`` query row and walks the whole - ``[0, context_length)`` KV window in ``BKV``-wide tiles. ``waves_per_block`` - waves each own a disjoint slice of the ``BKV/MFMA_N`` column tiles (no - cross-wave sharing / barrier), mirroring the dense kernel's wave split. - - ``kv_block_size`` (KVBlockSize) is the number of tokens co-packed in one - physical block; each lane gathers its own logical KV column independently - (``kv_indices[b, col // KVBlockSize]`` -> physical block, ``col % - KVBlockSize`` -> token within the block-flat block), so a column tile may - straddle block boundaries. ``KVBlockSize == 1`` degenerates to the per-token - ``[D fp8 | 4 scale]`` slot layout. - - ``preshuffle`` selects the production ``shuffle_weight(layout=(16,16))`` KV - data layout (each block's ``KVB*D`` fp8 bytes reordered into 16-token x - 16-hidden tiles; the co-packed f32 scale tail stays token-ordered). Only the - intra-block key-byte offset changes -- the block-table gather, scale gather, - MFMA, and mask are identical. Requires ``KVBlockSize % 16 == 0``. - - ``scale_mul`` folds the FN->FNUZ 2x-per-converted-operand compensation into - the (co-packed) per-token scale at compile time (ReLU positive-homogeneity). + """Paged MQA-logits kernel with SplitKV column parallelism. + + One thread block owns one ``(batch, next_n)`` query row and walks its KV + window in ``BKV``-wide tiles; the ``waves_per_block`` waves each take a + disjoint slice of the ``BKV/MFMA_N`` column tiles, with no barrier between + them. + + ``preshuffle`` only changes the intra-block key-byte offset -- the + block-table gather, scale gather, MFMA and mask are identical either way. + ``scale_mul`` folds the FN->FNUZ 2x compensation into the per-token scale at + compile time, which is exact by ReLU positive-homogeneity. """ H = num_heads D = head_size @@ -209,14 +192,11 @@ def kernel( lane_mod_N = umod(lane, MFMA_N) lane8 = lane_div_N * 8 - # fp8 operands loaded as vector<8xi32> per K-step (32 bytes = 4 i64s) - # for the 32x32x64 scaled MFMA. One vec_load_dwordx2 per i64 sub-step. q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) - # Uniform-base views over the co-packed cache: per-token gather is done - # via the (per-lane) byte offset, NOT the buffer base -- a per-lane base - # can't ride a scalar buffer descriptor. The byte offset stays i32 - # (assumes num_blocks*index_dim < 2^31; see host assert); a shared cache - # pool large enough to exceed that would need an i64/global-load gather. + # Uniform-base views: the per-token gather rides the per-lane byte + # offset, not the buffer base, since a per-lane base cannot ride a + # scalar buffer descriptor. That offset stays i32, hence the host's + # pool-size assert. kv_i32 = GTensor(KV_cache, dtype=T.i32, shape=(-1,)) kv_f32 = GTensor(KV_cache, dtype=T.f32, shape=(-1,)) w_t = GTensor(weights, dtype=T.f32, shape=(-1, H)) @@ -232,10 +212,9 @@ def kernel( out_row_t = make_out_row_view(out_logits, stride_out_i64, out_row) # ---- Preload Q frags + weights for the single query row ---- - # A-operand layout is per in-wave lane, so `lane` (not `tid`) indexes Q. - # With MFMA_K=64, each K-step loads a vector<8xi32> (4 packed i64s covering - # k = lane8 + [0,16,32,48]). fn_to_fnuz not supported for v8i32 operands; - # gfx950 uses native fp8 so convert_q_fn is always False on this arch. + # The A-operand layout is per in-wave lane, so `lane` indexes Q, not + # `tid`. convert_q_fn is always False here: the v8i32 operand has no + # fn_to_fnuz path, and gfx950 is native fp8 anyway. q_row_base = pid_batch * stride_q_batch + pid_next_n * stride_q_next_n a_pack = [[None] * K_STEPS for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): @@ -257,13 +236,10 @@ def kernel( h_w = mi * MFMA_M + h_off w_frag[mi][ii] = _to_raw(fx.Float32(w_t[out_row, h_w])) - # ---- SplitKV: this CTA owns a disjoint, BKV-aligned slice of the KV - # window. total_tiles = ceil(ctx/BKV); each split owns - # ceil(total_tiles/split_kv) contiguous tiles => slices are gap-free and - # disjoint, so every logits column has exactly one writer (pure - # parallelism across n, no cross-CTA reduction). split_kv==1 collapses to - # the full [0, ceil(ctx/BKV)*BKV) window. Splits past ctx early-exit - # (tile_lo_col >= tile_hi_col => the loop runs 0 iterations). ---- + # ---- SplitKV: each split owns ceil(total_tiles/split_kv) contiguous + # BKV-aligned tiles, so the slices are gap-free and disjoint and every + # logits column has exactly one writer -- no cross-CTA reduction. A + # split past ctx runs zero iterations. ---- context_chunk_num = arith.ceildivui( _to_raw(context_length), _to_raw(fx.Int32(BKV)) ) @@ -289,12 +265,9 @@ def kernel( col = col0 + abs_ni * MFMA_N + lane_mod_N # Clamp the gather index to a valid token (mask handles the rest). col_c = fx.Int32(arith.minsi(_to_raw(col), _to_raw(ctx_m1))) - # Block table maps a logical column to a physical block. The - # block-flat block byte-stride is KVB*index_dim, keys grouped - # first (tok*D) then all scales (KVB*D + tok*4). KVB is a - # compile-time constant, so for KVB==1 divui/remui/mul-by-1 fold - # away (block_idx==col_c, tok_in_block==0) -- no Python `if` here - # because the FlyDSL AST rewriter would turn it into a runtime + # KVB is a compile-time constant, so at KVB==1 the div/mod/mul + # fold away on their own. Deliberately no Python `if` for that + # case: the FlyDSL AST rewriter would turn it into a runtime # scf.if and drop the branch-local addresses. block_idx = udiv(col_c, KVB) tok_in_block = umod(col_c, KVB) @@ -302,12 +275,8 @@ def kernel( physical = fx.Int32(ind_t[ind_off]) block_byte_base = physical * block_byte_stride - # B operands: vector<8xi32> per K-step (MFMA_K=64). - # convert_kv_fn is always False (gfx950 native fp8). The - # preshuffle branch is compile-time (Python if): it only swaps - # the intra-block key-byte offset for the shuffle_weight(16x16) - # layout; the scale gather below is unchanged (scales stay - # co-packed and token-ordered). + # Compile-time branch: preshuffle only swaps the intra-block + # key-byte offset, leaving the scale gather below unchanged. b_col = [None] * K_STEPS if preshuffle: for kk in range_constexpr(K_STEPS): @@ -342,8 +311,7 @@ def kernel( ) # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. - # Write only in-causal columns; masked positions keep the - # caller's -inf prefill (matches clean_logits semantics). + # Masked positions keep the caller's -inf prefill. is_writer = _to_raw((lane_div_N == 0) & (col <= q_limit)) with ir.InsertionPoint(scf.IfOp(is_writer).then_block): out_row_t[col] = fx.Float32(col_sum) @@ -484,35 +452,26 @@ def flydsl_fp8_paged_mqa_logits( Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. - q_fp8: [batch, next_n, heads, hidden_dim], dtype float8 (e4m3 fn/fnuz) - kv_cache: [num_blocks, KVBlockSize, 1, index_dim], dtype uint8, block-flat - co-packed: per physical block, KVBlockSize*hidden_dim fp8 K - bytes (all key rows grouped) then KVBlockSize f32 dequant scales; - index_dim == hidden_dim + 4 (+ optional 16B padding). When - ``Preshuffle=True`` the fp8 data section is reordered by the - production ``shuffle_weight(layout=(16,16))`` (16-token x - 16-hidden tiles); the f32 scale tail stays token-ordered. - Preshuffle: consume the shuffle_weight(16x16) KV data layout. Requires - ``KVBlockSize % 16 == 0``. Default False (block-flat gather). - weights: [batch*next_n, heads], dtype float32 - out_logits: [batch*next_n, max_model_len], dtype float32. MUST be prefilled - with -inf by the caller; the kernel writes only in-window - (causal) positions and leaves masked positions untouched. - context_lens: [batch], dtype int32 - kv_indices: [batch, max_block_len], dtype int32 block table (physical block - per logical block ``p // KVBlockSize``; KVBlockSize==1 => a - per-token page table). - max_model_len: int, out_logits column count. + q_fp8: [batch, next_n, heads, hidden_dim], float8 e4m3 (fn or fnuz) + kv_cache: [num_blocks, KVBlockSize, 1, index_dim] uint8, co-packed per + block as KVBlockSize*hidden_dim fp8 K bytes then KVBlockSize + f32 scales; index_dim == hidden_dim + 4 (+ optional padding). + Preshuffle: consume the shuffle_weight(16x16) layout for the fp8 data + section (the scale tail stays token-ordered). Requires + ``KVBlockSize % 16 == 0``. + weights: [batch*next_n, heads], float32 + out_logits: [batch*next_n, max_model_len], float32. MUST be prefilled with + -inf: the kernel writes only in-window positions. + context_lens: [batch], int32 + kv_indices: [batch, max_block_len], int32 block table indexed by + ``p // KVBlockSize``. + max_model_len: out_logits column count. SplitKV: KV-column split count (grid = split_kv*batch*next_n). None => - sized so the dispatch lands WavePerEU waves on every SIMD; - pass an int to override, 1 disables splitting. Splits own - disjoint, gap-free BKV-aligned column ranges (one writer per - column, no cross-CTA reduction). - WavePerEU: target resident waves per SIMD for the auto-SplitKV formula - (default 2, which is where the measured curve bottoms out). + sized for WavePerEU waves per SIMD; 1 disables splitting. + WavePerEU: target resident waves per SIMD (default 2, where the measured + curve bottoms out). TotalCuCount: CU count for that formula; None => query the device. - ChunkK: KV tile width (must be a multiple of MFMA_N=32). Mirrors - the Triton ChunkK parameter. Defaults to 128. + ChunkK: KV tile width, a multiple of MFMA_N=32. Defaults to 128. Returns the same ``out_logits`` tensor (written in place). """ From 3de0f40053fe83e1621a24a26a717305abf3c3dd Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 15:31:35 +0000 Subject: [PATCH 14/52] refactor(flydsl): drive the paged mqa-logits tile loop with scf.for_ Replaces the hand-built scf.ForOp plus `with ir.InsertionPoint(body)` and trailing scf.YieldOp with the scf.for_ generator, which owns the insertion point and terminator. Same pattern the fmha_gfx1250 kernel already uses. Only two lines, but it removes the last raw MLIR region plumbing from the kernel's main loop. The masked store keeps scf.IfOp: FlyDSL exposes no `if_` helper, and switching to a predicated store would change codegen. ISA byte-identical (655 instructions), 870/870 correctness at 5.3e-4, 27.7us. Also investigated expressing the preshuffle as a FlyDSL layout, as examples/04-preshuffle_gemm.py does in two lines. Not adopted: that example indexes a whole contiguous tensor, whereas our access is a paged gather off a per-lane runtime base. GTensor._linear_offset is element-granular, and while every preshuffle stride happens to be dword-aligned, a 5-D view would just replace one multiply-add chain with another for about two lines. It only pays off alongside FlyDSL-native tensors and copy atoms. Co-authored-by: Cursor --- .../kernels/mqa_logits/fp8_paged_mqa_logits.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index ae6fbcea067..0c1cc8e7f81 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -253,12 +253,12 @@ def kernel( # loop-invariant, so hoist it out of the per-column gather below. block_byte_stride = fx.Int32(KVB) * index_dim - tile_lo = _to_raw(fx.Index(tile_lo_col)) - tile_hi = _to_raw(fx.Index(fx.Int32(tile_hi_col))) - tile_step = _to_raw(fx.Index(fx.Int32(BKV))) - tile_loop = scf.ForOp(tile_lo, tile_hi, tile_step, []) - with ir.InsertionPoint(tile_loop.body): - col0 = fx.Int32(arith.index_cast(T.i32, tile_loop.induction_variable)) + for iv in scf.for_( + _to_raw(fx.Index(tile_lo_col)), + _to_raw(fx.Index(fx.Int32(tile_hi_col))), + _to_raw(fx.Index(fx.Int32(BKV))), + ): + col0 = fx.Int32(arith.index_cast(T.i32, iv)) wave_ni_base = wave * N_TILES_PER_WAVE for ni in range_constexpr(N_TILES_PER_WAVE): abs_ni = wave_ni_base + ni @@ -317,8 +317,6 @@ def kernel( out_row_t[col] = fx.Float32(col_sum) scf.YieldOp([]) - scf.YieldOp([]) - @flyc.jit def launch_fp8_paged_mqa_logits( Q: fx.Tensor, From 90fcbf21e229f17eb9f5019c9117dd9a15847869 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 15:51:56 +0000 Subject: [PATCH 15/52] refactor(flydsl): express the KV preshuffle as a layout, not a code path The paged gather hand-computed byte offsets and branched on Preshuffle to pick between two of them. Both layouts are actually the same coordinate space -- (block, token//16, token%16, K-step, sub-group, lane half) -- and differ only in the stride vector: block-flat (blk, 4D, D/4, 16, 4, 2) token c at c*D shuffled (blk, 4D, 4, 256, 64, 2) token c at (c//16)*16D + (c%16)*16 So the choice moves into make_kv_key_view and the kernel loses its `if preshuffle` branch entirely; one load_pack_kv serves both. The co-packed scale tail becomes a (block, token) view whose base absorbs the KVB*D key offset, replacing the manual scale_byte arithmetic. load_pack_kv slices the view down to the `sub` axis once rather than indexing all six coordinates per load. Without that the invariant part of the address is re-derived in each of the four loads and the kernel grows 49 instructions; with it, 14. Trade, measured rather than assumed: the kernel file drops 15 lines and a branch, the shared module gains 21 for the two view constructors, so the pair is +6 overall. Codegen is 14 instructions larger (669 vs 655) because the hand-written version exploited disjoint bit-fields to fuse addresses into v_or3_b32/v_add3_u32 that the general dot product does not reach for. Speed holds on every path (3 reps each, layout vs previous): preshuffle 27.58 vs 27.83us, block-flat KVB=64 9.78 vs 9.98us, KVB=1 10.20 vs 10.53us. 870/870 correctness at 5.3e-4, no new ruff findings. Co-authored-by: Cursor --- .../kernels/mqa_logits/_mqa_logits_common.py | 61 +++++++++++++------ .../mqa_logits/fp8_paged_mqa_logits.py | 39 ++++-------- 2 files changed, 53 insertions(+), 47 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index 52b693c16a6..46f4a8f233e 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -94,31 +94,52 @@ def _load_i64(off): return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) -def load_pack_v8i32_preshuffle( - i32_view, block_byte_base, tok_in_block, head_size, kk_step, lane8 -): - """Same operand as ``load_pack_v8i32``, from the ``shuffle_weight(16,16)`` - layout the production path stores. Element (token c, hidden h) of a block - sits at ``(c//16)*(D*16) + (h//16)*256 + (c%16)*16 + (h%16)``, the inverse - of the host permute. Only the ``KVB*D`` data section is shuffled; the f32 - scale tail stays token-ordered, so the scale gather is unchanged. +def make_kv_key_view(kv_cache, head_size, kv_block_size, index_dim, preshuffle): + """Dword view of the co-packed KV keys over the gather's own axes: + ``(block, token//16, token%16, K-step, sub-group, lane half)``. + + The two data layouts differ only in this stride vector -- block-flat puts + token c at ``c*D``, shuffle_weight(16,16) at ``(c//16)*16D + (c%16)*16`` + with hidden groups every 256B -- so Preshuffle is a layout, not a second + gather path. index_dim is a multiple of 4, so the block stride divides. """ D = head_size - # Loop-invariant across the four hidden groups, so hoisted out of the load. - c_hi = udiv(tok_in_block, 16) - c_lo = umod(tok_in_block, 16) - base = block_byte_base + c_hi * (D * 16) + c_lo * 16 + lane8 + blk = udiv(fx.Int32(kv_block_size) * index_dim, 4) + stride = ( + (blk, D * 4, 4, 256, 64, 2) if preshuffle else (blk, D * 4, D // 4, 16, 4, 2) + ) + return GTensor(kv_cache, dtype=T.i32, shape=(-1,) * 6, stride=stride) + + +def make_kv_scale_view(kv_cache, head_size, kv_block_size, index_dim): + """f32 ``(block, token)`` view of the co-packed scales. The tail follows the + KVB*D key bytes and is token-ordered in both layouts, so it folds into base. + """ + return GTensor( + kv_cache, + dtype=T.f32, + shape=(-1, kv_block_size), + stride=(udiv(fx.Int32(kv_block_size) * index_dim, 4), 1), + base_offset=kv_block_size * head_size // 4, + ) + - def _load_i64(g): - v2 = i32_view.vec_load((udiv(base + g * 256, 4),), vec_size=2) +def load_pack_kv(key_view, physical, tok_in_block, kk_step, lane_div_N): + """vector<8xi32> B operand for one KV token and K-step. Layout-agnostic: + the view's strides carry the block-flat/preshuffle difference. + """ + c_hi, c_lo = udiv(tok_in_block, 16), umod(tok_in_block, 16) + # Slice everything but `sub` first: that fixes the token's base offset once + # instead of re-deriving it inside each of the four loads. + sub_view = key_view[physical, c_hi, c_lo, kk_step, None, lane_div_N] + + def _load_i64(sub): + v2 = sub_view.vec_load((sub,), vec_size=2) return Vec(v2).bitcast(fx.Int64)[0].ir_value() - g0 = kk_step * 4 # first hidden-16 group of this MFMA K-step - i64_0 = _load_i64(g0 + 0) - i64_1 = _load_i64(g0 + 1) - i64_2 = _load_i64(g0 + 2) - i64_3 = _load_i64(g0 + 3) - return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) + return Vec.from_elements( + [_load_i64(sub) for sub in range_constexpr(4)], fx.Int64 + ).bitcast(fx.Int32) def make_out_row_view(logits, stride_i64, row_i32): diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 0c1cc8e7f81..cb630e420d8 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -54,8 +54,10 @@ DEFAULT_COMPILE_HINTS, Vec, device_cu_count, + load_pack_kv, load_pack_v8i32, - load_pack_v8i32_preshuffle, + make_kv_key_view, + make_kv_scale_view, make_out_row_view, mfma_head_reduce, udiv, @@ -197,8 +199,8 @@ def kernel( # offset, not the buffer base, since a per-lane base cannot ride a # scalar buffer descriptor. That offset stays i32, hence the host's # pool-size assert. - kv_i32 = GTensor(KV_cache, dtype=T.i32, shape=(-1,)) - kv_f32 = GTensor(KV_cache, dtype=T.f32, shape=(-1,)) + kv_key = make_kv_key_view(KV_cache, D, KVB, index_dim, preshuffle) + kv_scale_v = make_kv_scale_view(KV_cache, D, KVB, index_dim) w_t = GTensor(weights, dtype=T.f32, shape=(-1, H)) cl_t = GTensor(context_lens, dtype=T.i32, shape=(-1,)) ind_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) @@ -249,9 +251,6 @@ def kernel( tile_lo_col = pid_split_kv * split_cols tile_hi_col = arith.minsi(_to_raw(tile_lo_col + split_cols), _to_raw(full_end)) ctx_m1 = context_length - 1 - # Block-flat block byte-stride (KVB tokens per physical block); - # loop-invariant, so hoist it out of the per-column gather below. - block_byte_stride = fx.Int32(KVB) * index_dim for iv in scf.for_( _to_raw(fx.Index(tile_lo_col)), @@ -269,29 +268,15 @@ def kernel( # fold away on their own. Deliberately no Python `if` for that # case: the FlyDSL AST rewriter would turn it into a runtime # scf.if and drop the branch-local addresses. - block_idx = udiv(col_c, KVB) tok_in_block = umod(col_c, KVB) - ind_off = pid_batch * max_block_len + block_idx + ind_off = pid_batch * max_block_len + udiv(col_c, KVB) physical = fx.Int32(ind_t[ind_off]) - block_byte_base = physical * block_byte_stride - - # Compile-time branch: preshuffle only swaps the intra-block - # key-byte offset, leaving the scale gather below unchanged. - b_col = [None] * K_STEPS - if preshuffle: - for kk in range_constexpr(K_STEPS): - b_col[kk] = load_pack_v8i32_preshuffle( - kv_i32, block_byte_base, tok_in_block, D, kk, lane8 - ) - else: - key_byte = block_byte_base + tok_in_block * D - for kk in range_constexpr(K_STEPS): - byte_base_b = key_byte + kk * MFMA_K - b_col[kk] = load_pack_v8i32(kv_i32, byte_base_b, lane8) - - # Co-packed per-token f32 scale (block-flat: after the KVB keys). - scale_byte = block_byte_base + (KVB * D) + tok_in_block * 4 - kv_scale = _to_raw(fx.Float32(kv_f32[udiv(scale_byte, 4)])) + + b_col = [ + load_pack_kv(kv_key, physical, tok_in_block, kk, lane_div_N) + for kk in range_constexpr(K_STEPS) + ] + kv_scale = _to_raw(fx.Float32(kv_scale_v[physical, tok_in_block])) if scale_mul != 1.0: kv_scale = arith.MulFOp( kv_scale, scale_mul_c, fastmath=fm_fast From 7ae5ec5268aad481d1eb31ff73482ee5991f9d13 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 22:36:59 +0000 Subject: [PATCH 16/52] test(flydsl): curate the paged mqa-logits matrix, 870 cases -> 29 The sweep crossed seven axes into 864 cases that compiled only 12 distinct kernels, re-running each about 72 times. The kernel specialises at compile time on (heads, head_dim, KVBlockSize, Preshuffle); batch, next_n, context length, SplitKV and the Q dtype are runtime values, and none of the axes interact -- nothing couples head_dim to SplitKV or the dtype to KVBlockSize -- so the product was testing interactions that cannot exist. q_dtype alone cost 432 cases while selecting no different code path on gfx950, where convert_q_fn is always False. Now: a full cross of the compile-time axes (16 cases, one per distinct kernel -- more preshuffle coverage than before, which only hand-picked 4), each runtime axis varied alone off a base, and four interaction cases where the tile-range arithmetic genuinely couples. --exhaustive still runs the original 870 for release validation. Validated by mutation testing rather than by the suite passing, since a suite cut too far also passes. Eleven deliberate kernel bugs -- swapped preshuffle and block-flat strides, causal off-by-one, dropped tail clamp, wrong lane half, wrong block-table index, dropped kv_scale, wrong head weight, floor-instead-of-ceil split (needs SplitKV>1), batch-blind output row (needs batch>1), next_n-blind causal bound (needs next_n>1). The full suite catches 11/11; the curated matrix also catches 11/11. 29 cases in 8.7s against 870 in 79s. Harness in att_profiles/mutation_check.py. Co-authored-by: Cursor --- .../test_flydsl_fp8_paged_mqa_logits.py | 179 +++++++++++------- 1 file changed, 109 insertions(+), 70 deletions(-) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 2cd38c013d8..ecbacd71d40 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -313,20 +313,105 @@ def test_fp8_paged_mqa_logits( } -# Preshuffle is a distinct correctness branch (production shuffle_weight(16x16) -# data layout, KVBlockSize % 16 == 0). A small fixed set covers what differs from -# the block-flat path -- KVBlockSize 16 and 64, head_dim 64 and 128, the FN->FNUZ -# Q patch, and a couple of grid/context shapes -- without re-running the whole -# block-flat sweep. -# (batch_size, next_n, heads, head_dim, avg_kv, q_dtype, kv_block_size) -PRESHUFFLE_CASES = [ - (1, 1, 64, 128, 1024, "fnuz", 16), - (1, 2, 64, 128, 1024, "fnuz", 64), - (2, 1, 128, 128, 8192, "fnuz", 64), - (1, 1, 64, 64, 1024, "fnuz", 16), # head_dim=64 (single K-step) - (1, 1, 64, 128, 1024, "fn", 64), # exercises the gfx942 FN->FNUZ Q patch - (4, 2, 64, 128, 8192, "fnuz", 64), -] +# The kernel specialises at compile time on (heads, head_dim, KVBlockSize, +# Preshuffle); batch, next_n, context length, SplitKV and the Q dtype are +# runtime values. Those axes do not interact -- nothing couples head_dim to +# SplitKV, or the dtype to KVBlockSize -- so crossing all seven produced 864 +# cases that compiled only 12 distinct kernels and re-ran each about 72 times. +# +# Instead: cross the compile-time axes, vary each runtime axis alone from a +# base, and add back the few combinations where the tile-range arithmetic +# genuinely couples. `--exhaustive` still runs the full product. +_BASE = { + "batch_size": 2, + "next_n": 2, + "heads": 64, + "head_dim": 128, + "avg_kv_length": 1024, + "q_dtype": "fnuz", +} +_SHAPES = [(64, 64), (64, 128), (128, 64), (128, 128)] + + +def default_cases(): + """Curated case list; each entry is a kwargs dict for the test body.""" + cases = [] + # Compile-time cross: one entry per distinct compiled kernel. + for heads, head_dim in _SHAPES: + for kvb in (1, 64): + cases.append({**_BASE, "heads": heads, "head_dim": head_dim, "block_size": kvb}) + for kvb in (16, 64): # Preshuffle needs KVBlockSize % 16 == 0 + cases.append( + { + **_BASE, + "heads": heads, + "head_dim": head_dim, + "block_size": kvb, + "preshuffle": True, + } + ) + # Runtime axes, one at a time off the base. + for bs, nn in ((1, 1), (1, 2), (4, 2), (8, 1)): + cases.append({**_BASE, "batch_size": bs, "next_n": nn, "block_size": 64}) + for kv in (128, 8192): # 128 is shorter than one ChunkK tile + cases.append({**_BASE, "avg_kv_length": kv, "block_size": 64}) + for sk in (1, 4): # 1 disables splitting; base uses the auto formula + cases.append({**_BASE, "split_kv": sk, "block_size": 64}) + cases.append({**_BASE, "q_dtype": "fn", "block_size": 64}) # FN->FNUZ Q patch + # Interactions: SplitKV against the context length it partitions, and + # against the preshuffled layout (which the old matrix never crossed). + cases.append({**_BASE, "avg_kv_length": 128, "split_kv": 4, "block_size": 64}) + cases.append({**_BASE, "avg_kv_length": 8192, "split_kv": 1, "block_size": 64}) + cases.append({**_BASE, "split_kv": 4, "block_size": 64, "preshuffle": True}) + cases.append({**_BASE, "split_kv": 1, "block_size": 16, "preshuffle": True}) + return cases + + +def exhaustive_cases(): + """The original 7-way cartesian product, for release validation.""" + cases = [] + for (bs, nn), nh, hd, kv, qd, sk, kvb in itertools.product( + [(1, 1), (1, 2), (2, 1), (2, 2), (4, 2), (8, 1)], + [64, 128], + [64, 128], + [128, 1024, 8192], + ["fnuz", "fn"], + [0, 1, 4], + [1, 64], + ): + cases.append( + { + "batch_size": bs, + "next_n": nn, + "heads": nh, + "head_dim": hd, + "avg_kv_length": kv, + "q_dtype": qd, + "split_kv": sk, + "block_size": kvb, + } + ) + for bs, nn, nh, hd, kv, qd, kvb in [ + (1, 1, 64, 128, 1024, "fnuz", 16), + (1, 2, 64, 128, 1024, "fnuz", 64), + (2, 1, 128, 128, 8192, "fnuz", 64), + (1, 1, 64, 64, 1024, "fnuz", 16), + (1, 1, 64, 128, 1024, "fn", 64), + (4, 2, 64, 128, 8192, "fnuz", 64), + ]: + cases.append( + { + "batch_size": bs, + "next_n": nn, + "heads": nh, + "head_dim": hd, + "avg_kv_length": kv, + "q_dtype": qd, + "block_size": kvb, + "preshuffle": True, + } + ) + return cases def main(): @@ -340,70 +425,24 @@ def main(): return parser = argparse.ArgumentParser( - formatter_class=argparse.RawTextHelpFormatter, - description="FlyDSL paged fp8_mqa_logits correctness sweep", - ) - parser.add_argument( - "-b", - "--batch-next-n", - type=dtypes.str2tuple, - nargs="*", - default=[(1, 1), (1, 2), (2, 1), (2, 2), (4, 2), (8, 1)], - help="(batch_size, next_n) pairs", - ) - parser.add_argument("--num-heads", type=int, nargs="*", default=[64, 128]) - parser.add_argument("--head-dim", type=int, nargs="*", default=[64, 128]) - parser.add_argument( - "--avg-kv-length", type=int, nargs="*", default=[128, 1024, 8192] + description="FlyDSL paged fp8_mqa_logits correctness sweep" ) parser.add_argument( - "--q-dtype", - type=str, - nargs="*", - default=["fnuz", "fn"], - choices=["fnuz", "fn"], - ) - parser.add_argument( - "--split-kv", - type=int, - nargs="*", - default=[0, 1, 4], - help="0 == auto (production formula); else explicit SplitKV (1 disables)", + "--exhaustive", + action="store_true", + help="run the full 7-way cartesian product (870 cases) instead of the " + "curated matrix; for release validation", ) parser.add_argument( - "--kv-block-size", - type=int, - nargs="*", - default=[1, 64], - help="KVBlockSize (paged block-flat page size); 1 == per-token slots", - ) - parser.add_argument( - "--no-preshuffle", - action="store_true", - help="skip the preshuffle (shuffle_weight 16x16 layout) correctness cases", + "--no-preshuffle", action="store_true", help="skip the preshuffle cases" ) args = parser.parse_args() - df = [] - for (bs, nn), nh, hd, kv, qd, sk, kvb in itertools.product( - args.batch_next_n, - args.num_heads, - args.head_dim, - args.avg_kv_length, - args.q_dtype, - args.split_kv, - args.kv_block_size, - ): - df.append(test_fp8_paged_mqa_logits(bs, nn, nh, hd, kv, qd, sk, block_size=kvb)) - - if not args.no_preshuffle: - for bs, nn, nh, hd, kv, qd, kvb in PRESHUFFLE_CASES: - df.append( - test_fp8_paged_mqa_logits( - bs, nn, nh, hd, kv, qd, 0, block_size=kvb, preshuffle=True - ) - ) + cases = exhaustive_cases() if args.exhaustive else default_cases() + if args.no_preshuffle: + cases = [c for c in cases if not c.get("preshuffle")] + df = [test_fp8_paged_mqa_logits(**c) for c in cases] df = pd.DataFrame(df) try: summary = df.to_markdown(index=False) From 8979a091749d9815300eb3d0861f4c4a4313869c Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 22:42:01 +0000 Subject: [PATCH 17/52] test(flydsl): cover the wave-split and ChunkK axes Neither was ever exercised: every case ran the default variant at ChunkK=128, so paged_w2 and paged_w4 had no coverage at all, and the Tier-1 change of DEFAULT_VARIANT silently moved which one CI touched. That hole was real, not theoretical. A mutation of the wave-to-column-tile assignment (`wave * N_TILES_PER_WAVE` -> `wave * (N_TILES_PER_WAVE - 1)`) is correct by construction at one wave per CTA, where `wave` is always 0, and wrong for any multi-wave variant. It survived both the curated suite *and* the original 870-case exhaustive sweep. Adds variant and chunk_k parameters plumbed through to the kernel, reported in the results table, and six cases: paged_w2 and paged_w4 at the default tile width, ChunkK 64 and 256 against the variants that divide them, and a preshuffle x paged_w4 cross. 35 cases in 9.0s, and the mutation set now runs 12/12 where the exhaustive suite manages 11/12. Co-authored-by: Cursor --- .../test_flydsl_fp8_paged_mqa_logits.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index ecbacd71d40..96d3360cb42 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -224,6 +224,8 @@ def test_fp8_paged_mqa_logits( split_kv=0, block_size=1, preshuffle=False, + variant=None, + chunk_k=128, ): # split_kv == 0 -> auto (production host formula); else an explicit override # (1 disables splitting). Both must be correctness-identical. @@ -285,6 +287,8 @@ def test_fp8_paged_mqa_logits( Preshuffle=preshuffle, KVBlockSize=block_size, SplitKV=_split_kv, + ChunkK=chunk_k, + variant=variant, ) got_mask = got == float("-inf") @@ -309,6 +313,8 @@ def test_fp8_paged_mqa_logits( "kvb": block_size, "preshuffle": preshuffle, "split_kv": "auto" if split_kv == 0 else split_kv, + "variant": variant or "default", + "chunk_k": chunk_k, "flydsl err": err, } @@ -364,6 +370,18 @@ def default_cases(): cases.append({**_BASE, "avg_kv_length": 8192, "split_kv": 1, "block_size": 64}) cases.append({**_BASE, "split_kv": 4, "block_size": 64, "preshuffle": True}) cases.append({**_BASE, "split_kv": 1, "block_size": 16, "preshuffle": True}) + # Wave split and KV tile width. Neither was covered before -- every case ran + # the default variant at ChunkK=128 -- so a bug in the wave-to-column-tile + # assignment was invisible, being correct by construction at one wave per + # CTA. ChunkK/32 must be divisible by the variant's wave count. + for variant in ("paged_w2", "paged_w4"): + cases.append({**_BASE, "block_size": 64, "variant": variant}) + cases.append({**_BASE, "block_size": 64, "chunk_k": 64, "variant": "paged_w2"}) + cases.append({**_BASE, "block_size": 64, "chunk_k": 256, "variant": "paged_w4"}) + cases.append({**_BASE, "block_size": 64, "chunk_k": 256}) + cases.append( + {**_BASE, "block_size": 64, "preshuffle": True, "variant": "paged_w4"} + ) return cases From 490156e2e65c789f4591a7e3edcc2bc404d36b3b Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 14 Aug 2026 22:47:45 +0000 Subject: [PATCH 18/52] test(flydsl): condense the paged mqa-logits test prose and plumbing Docstrings trimmed to what a reader cannot derive -- the oracle's formula and mask rule, the co-pack byte layout, what the preshuffled copy is for -- dropping the restatement and the cross-references to other files. _build_inputs returned an 8-tuple that was unpacked over ten lines at the call site; it now returns a NamedTuple, so the test body reads inp.q_fp8 rather than positional unpacking, and the contract is documented by the type. Its block table is also built vectorized instead of one element per GPU write, which is why the suite went from 9.0s to 7.5s. Same values: the per-case errors are unchanged to the digit. The pool there is still sized from max_model_len, so sequences share blocks once the batch outgrows it. Left alone deliberately, with a comment saying why: harmless for correctness and arguably extra coverage of block reuse, unlike in the benchmarks where it shrank the measured working set into cache. Still 35 cases, still 12/12 on the mutation set. Co-authored-by: Cursor --- .../test_flydsl_fp8_paged_mqa_logits.py | 254 ++++++++---------- 1 file changed, 118 insertions(+), 136 deletions(-) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 96d3360cb42..365a04dabe6 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -3,17 +3,15 @@ """Correctness sweep for the FlyDSL paged FP8 MQA-logits (decode) kernel. -Mirrors the aiter-op-test standard used by ``test_flydsl_fp8_mqa_logits.py`` and -the paged benchmark ``op_tests/op_benchmarks/triton/bench_deepgemm_attention.py``. -The correctness reference is a torch port of vLLM's ``fp8_paged_mqa_logits_torch`` -(dequantizes the co-packed fp8 cache and applies the causal mask); it is copied -here so the test does not depend on vLLM. Correctness gate: exact ``-inf``-mask -match + ``calc_diff < 1e-3`` (tolerances are NOT widened). +The reference is a torch port of vLLM's ``fp8_paged_mqa_logits_torch``, copied +here so the test does not depend on vLLM. Gate: exact ``-inf``-mask match plus +``calc_diff < 1e-3``, tolerances NOT widened. """ import argparse import itertools import random +from typing import NamedTuple import aiter import pandas as pd @@ -46,11 +44,8 @@ def calc_diff(x, y): def kv_cache_cast_to_fp8(x, fp8_dtype): - """Co-pack a bf16 KV cache into the fp8+scale byte layout (KVBlockSize>=1). - - Layout per block-row: KVBlockSize*head_dim fp8 bytes, then KVBlockSize f32 - (4-byte) per-token scales. No 16B padding (Phase 1). Mirrors the paged - benchmark's builder. + """Co-pack a bf16 KV cache into the fp8+scale byte layout: per block-row, + KVBlockSize*head_dim fp8 bytes then KVBlockSize f32 scales, no padding. """ num_blocks, block_size, num_heads, head_dim = x.shape assert num_heads == 1 @@ -72,14 +67,10 @@ def kv_cache_cast_to_fp8(x, fp8_dtype): def preshuffle_kv_data(kv_cache_fp8, head_dim): - """Reorder the fp8 data section into the production ``Preshuffle`` layout. - - Applies ``shuffle_weight(layout=(16,16))`` to the per-block ``KVB*head_dim`` - fp8 key bytes (reshaped ``[num_blocks, KVBlockSize, head_dim]``) and leaves - the co-packed f32 scale tail untouched -- exactly what the production - ``deepgemm_fp8_paged_mqa_logits`` benchmark builds for ``Preshuffle=True``. - The torch oracle still reads the *unshuffled* cache (the reference is - layout-agnostic), so this shuffled copy is fed only to the kernel. + """Apply ``shuffle_weight(layout=(16,16))`` to the per-block fp8 key bytes, + leaving the co-packed f32 scale tail alone -- the production Preshuffle + layout. Only the kernel gets this copy; the oracle reads the unshuffled + cache, since the reference is layout-agnostic. """ from aiter.ops.shuffle import shuffle_weight @@ -108,14 +99,12 @@ def ref_fp8_paged_mqa_logits( fp8_dtype, block_size=1, ): - """Torch reference (vectorized port of vLLM ``fp8_paged_mqa_logits_torch``). + """Torch reference, vectorized port of vLLM ``fp8_paged_mqa_logits_torch``:: - Dequantizes the **block-flat** co-packed fp8 cache and computes: logits[b*next_n+n, p] = sum_h ReLU() * weights[.., h] - with the causal mask ``p <= context_len - next_n + n``. Logical position ``p`` - resolves to physical ``(block, tok) = (block_tables[b, p//KVB], p%KVB)`` and - the block-flat block groups all ``KVB`` fp8 key rows first, then all ``KVB`` - f32 scales. ``block_size == 1`` degenerates to the per-token page table. + + masked by ``p <= context_len - next_n + n``. Position ``p`` resolves to + ``(block_tables[b, p//KVB], p%KVB)`` in the block-flat co-packed cache. """ batch_size, next_n, heads, dim = q.size() num_blocks = kv_cache_fp8.shape[0] @@ -165,6 +154,17 @@ def ref_fp8_paged_mqa_logits( return logits +class Inputs(NamedTuple): + q: torch.Tensor # bf16, for the oracle + q_fp8: torch.Tensor # quantized, for the kernel + kv_cache_fp8: torch.Tensor # co-packed, unshuffled + weights: torch.Tensor + context_lens: torch.Tensor + block_tables: torch.Tensor + max_model_len: int + fp8_dtype: torch.dtype + + def _build_inputs( batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, block_size=1, seed=0 ): @@ -185,23 +185,27 @@ def _build_inputs( kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) weights = torch.randn((batch_size * next_n, heads), dtype=torch.float32) - # block table: one entry per KVBlockSize-token block, ceil(ctx / block_size). - max_block_len = (int(context_lens.max().item()) + block_size - 1) // block_size - block_tables = torch.zeros( - (batch_size, max_block_len), device="cuda", dtype=torch.int32 - ) + # Block table: one entry per KVBlockSize-token block, handed out of a + # shuffled pool in sequence order. The pool is sized from max_model_len, so + # sequences share blocks once the batch needs more than it holds -- harmless + # here (and extra coverage of block reuse), unlike in the benchmarks where + # it would shrink the measured working set into cache. + blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size + max_block_len = int(blocks_per_seq.max().item()) pool = list(range(num_blocks)) random.shuffle(pool) - counter = 0 - for i in range(batch_size): - n_blk = (int(context_lens[i].item()) + block_size - 1) // block_size - for j in range(n_blk): - block_tables[i][j] = pool[counter % num_blocks] - counter += 1 + pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) + col = torch.arange(max_block_len, device="cuda", dtype=torch.int64) + starts = torch.cumsum(blocks_per_seq, 0) - blocks_per_seq + block_tables = torch.where( + col[None, :] < blocks_per_seq[:, None], + pool_t[(starts[:, None] + col[None, :]) % num_blocks], + torch.zeros((), device="cuda", dtype=torch.int32), + ).to(torch.int32) q_fp8 = q.to(q_dtype) kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, fp8_dtype) - return ( + return Inputs( q, q_fp8, kv_cache_fp8, @@ -227,19 +231,10 @@ def test_fp8_paged_mqa_logits( variant=None, chunk_k=128, ): - # split_kv == 0 -> auto (production host formula); else an explicit override - # (1 disables splitting). Both must be correctness-identical. + # split_kv == 0 -> auto; else an explicit override (1 disables splitting). + # Both must be correctness-identical. _split_kv = None if split_kv == 0 else split_kv - ( - q, - q_fp8, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - ) = _build_inputs( + inp = _build_inputs( batch_size, next_n, heads, @@ -251,39 +246,40 @@ def test_fp8_paged_mqa_logits( with torch.inference_mode(): ref = ref_fp8_paged_mqa_logits( - q, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, + inp.q, + inp.kv_cache_fp8, + inp.weights, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + inp.fp8_dtype, block_size=block_size, ) ref_mask = ref == float("-inf") out = torch.full( - (batch_size * next_n, max_model_len), + (batch_size * next_n, inp.max_model_len), float("-inf"), device="cuda", dtype=torch.float32, ) - # Oracle reads the unshuffled cache above; the kernel gets the preshuffled - # copy (same numbers, production shuffle_weight(16x16) data layout). + # Same numbers as the oracle's cache, in the production data layout. kv_cache_kernel = ( - preshuffle_kv_data(kv_cache_fp8, head_dim) if preshuffle else kv_cache_fp8 + preshuffle_kv_data(inp.kv_cache_fp8, head_dim) + if preshuffle + else inp.kv_cache_fp8 ) with torch.inference_mode(): got = flydsl_fp8_paged_mqa_logits( - q_fp8, + inp.q_fp8, kv_cache_kernel, - weights, + inp.weights, out, - context_lens, - block_tables, - max_model_len, + inp.context_lens, + inp.block_tables, + inp.max_model_len, Preshuffle=preshuffle, KVBlockSize=block_size, SplitKV=_split_kv, @@ -339,56 +335,59 @@ def test_fp8_paged_mqa_logits( _SHAPES = [(64, 64), (64, 128), (128, 64), (128, 128)] +_AXES = ( + "batch_size", + "next_n", + "heads", + "head_dim", + "avg_kv_length", + "q_dtype", + "split_kv", + "block_size", +) + + +def _c(**kw): + """A case: the base config with KVBlockSize=64, plus the overrides given.""" + return {**_BASE, "block_size": 64, **kw} + + def default_cases(): - """Curated case list; each entry is a kwargs dict for the test body.""" cases = [] - # Compile-time cross: one entry per distinct compiled kernel. + # Compile-time cross: one case per distinct compiled kernel. Preshuffle + # needs KVBlockSize % 16 == 0. for heads, head_dim in _SHAPES: for kvb in (1, 64): - cases.append({**_BASE, "heads": heads, "head_dim": head_dim, "block_size": kvb}) - for kvb in (16, 64): # Preshuffle needs KVBlockSize % 16 == 0 + cases.append(_c(heads=heads, head_dim=head_dim, block_size=kvb)) + for kvb in (16, 64): cases.append( - { - **_BASE, - "heads": heads, - "head_dim": head_dim, - "block_size": kvb, - "preshuffle": True, - } + _c(heads=heads, head_dim=head_dim, block_size=kvb, preshuffle=True) ) # Runtime axes, one at a time off the base. - for bs, nn in ((1, 1), (1, 2), (4, 2), (8, 1)): - cases.append({**_BASE, "batch_size": bs, "next_n": nn, "block_size": 64}) - for kv in (128, 8192): # 128 is shorter than one ChunkK tile - cases.append({**_BASE, "avg_kv_length": kv, "block_size": 64}) - for sk in (1, 4): # 1 disables splitting; base uses the auto formula - cases.append({**_BASE, "split_kv": sk, "block_size": 64}) - cases.append({**_BASE, "q_dtype": "fn", "block_size": 64}) # FN->FNUZ Q patch + cases += [_c(batch_size=b, next_n=n) for b, n in ((1, 1), (1, 2), (4, 2), (8, 1))] + cases += [_c(avg_kv_length=kv) for kv in (128, 8192)] # 128 < one ChunkK tile + cases += [_c(split_kv=sk) for sk in (1, 4)] # 1 disables splitting + cases.append(_c(q_dtype="fn")) # gfx942 FN->FNUZ Q patch # Interactions: SplitKV against the context length it partitions, and - # against the preshuffled layout (which the old matrix never crossed). - cases.append({**_BASE, "avg_kv_length": 128, "split_kv": 4, "block_size": 64}) - cases.append({**_BASE, "avg_kv_length": 8192, "split_kv": 1, "block_size": 64}) - cases.append({**_BASE, "split_kv": 4, "block_size": 64, "preshuffle": True}) - cases.append({**_BASE, "split_kv": 1, "block_size": 16, "preshuffle": True}) - # Wave split and KV tile width. Neither was covered before -- every case ran - # the default variant at ChunkK=128 -- so a bug in the wave-to-column-tile - # assignment was invisible, being correct by construction at one wave per + # against the preshuffled layout, which the old matrix never crossed. + cases.append(_c(avg_kv_length=128, split_kv=4)) + cases.append(_c(avg_kv_length=8192, split_kv=1)) + cases.append(_c(split_kv=4, preshuffle=True)) + cases.append(_c(split_kv=1, block_size=16, preshuffle=True)) + # Wave split and KV tile width, neither covered before: a bug in the + # wave-to-column-tile assignment is correct by construction at one wave per # CTA. ChunkK/32 must be divisible by the variant's wave count. - for variant in ("paged_w2", "paged_w4"): - cases.append({**_BASE, "block_size": 64, "variant": variant}) - cases.append({**_BASE, "block_size": 64, "chunk_k": 64, "variant": "paged_w2"}) - cases.append({**_BASE, "block_size": 64, "chunk_k": 256, "variant": "paged_w4"}) - cases.append({**_BASE, "block_size": 64, "chunk_k": 256}) - cases.append( - {**_BASE, "block_size": 64, "preshuffle": True, "variant": "paged_w4"} - ) + cases += [_c(variant=v) for v in ("paged_w2", "paged_w4")] + cases.append(_c(chunk_k=64, variant="paged_w2")) + cases.append(_c(chunk_k=256, variant="paged_w4")) + cases.append(_c(chunk_k=256)) + cases.append(_c(preshuffle=True, variant="paged_w4")) return cases def exhaustive_cases(): - """The original 7-way cartesian product, for release validation.""" - cases = [] - for (bs, nn), nh, hd, kv, qd, sk, kvb in itertools.product( + """The original 7-way cartesian product (870 cases), for release validation.""" + prod = itertools.product( [(1, 1), (1, 2), (2, 1), (2, 2), (4, 2), (8, 1)], [64, 128], [64, 128], @@ -396,40 +395,23 @@ def exhaustive_cases(): ["fnuz", "fn"], [0, 1, 4], [1, 64], - ): - cases.append( - { - "batch_size": bs, - "next_n": nn, - "heads": nh, - "head_dim": hd, - "avg_kv_length": kv, - "q_dtype": qd, - "split_kv": sk, - "block_size": kvb, - } - ) - for bs, nn, nh, hd, kv, qd, kvb in [ - (1, 1, 64, 128, 1024, "fnuz", 16), - (1, 2, 64, 128, 1024, "fnuz", 64), - (2, 1, 128, 128, 8192, "fnuz", 64), - (1, 1, 64, 64, 1024, "fnuz", 16), - (1, 1, 64, 128, 1024, "fn", 64), - (4, 2, 64, 128, 8192, "fnuz", 64), - ]: - cases.append( - { - "batch_size": bs, - "next_n": nn, - "heads": nh, - "head_dim": hd, - "avg_kv_length": kv, - "q_dtype": qd, - "block_size": kvb, - "preshuffle": True, - } - ) - return cases + ) + cases = [ + dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, sk, kvb))) + for (bs, nn), nh, hd, kv, qd, sk, kvb in prod + ] + # Preshuffle was covered by a hand-picked set rather than a full cross. + return cases + [ + {**dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, 0, kvb))), "preshuffle": True} + for bs, nn, nh, hd, kv, qd, kvb in [ + (1, 1, 64, 128, 1024, "fnuz", 16), + (1, 2, 64, 128, 1024, "fnuz", 64), + (2, 1, 128, 128, 8192, "fnuz", 64), + (1, 1, 64, 64, 1024, "fnuz", 16), + (1, 1, 64, 128, 1024, "fn", 64), + (4, 2, 64, 128, 8192, "fnuz", 64), + ] + ] def main(): From 9a533274e4dd082bd58e44c9e391b9971005b906 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Sat, 15 Aug 2026 07:25:22 +0000 Subject: [PATCH 19/52] docs(flydsl): drop benchmark figures from the paged mqa-logits source Timings and percentages go stale as soon as anything around them changes, and the source is the wrong place to keep them -- the numbers behind these choices are in the commit history and the PR description. Removes the measured wave-count curve from _auto_split_kv, the "4-14% slower" claim from the DEFAULT_VARIANT note, and the "where the measured curve bottoms out" aside from the WavePerEU docs, keeping what the formula computes and why resident wave count is the lever. In the test, the block table comment keeps its warning that shared blocks are deliberate but no longer cross-references what that would cost a benchmark. Comments only; 35/35 correctness cases unchanged at 3.97e-4. Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits.py | 27 +++++++------------ .../test_flydsl_fp8_paged_mqa_logits.py | 5 ++-- 2 files changed, 11 insertions(+), 21 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index cb630e420d8..f4bf9121345 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -75,17 +75,10 @@ def _auto_split_kv( SplitKV = round(WavePerEU * 4 * TotalCuCount / (batch*next_n*WavePerBlock)) - Resident wave count is the lever that matters here -- the kernel is bound by - memory throughput, not per-wave latency, so too few waves leaves the memory - system idle and too many just queue. Measured on gfx950/256 CU across six - shapes (batch 4..128, kv_len 4k..128k), the curve is flat between 2.0 and 2.5 - waves/SIMD and rises steeply either side: at the reference shape 1.0/SIMD - costs 33.9 us and 4.0/SIMD costs 29.4 us against 27.6 us at 2.0. - - The earlier formula, ``((TotalCuCount // (batch*next_n) + 4) // 5 * 5) * - WavePerEU``, ignored the CTA's wave count, so it drifted with shape -- - 2.5 waves/SIMD at the reference shape but 10 at batch=128 -- and left 4-14% - on the table. Returns >= 1. + Resident wave count is the lever that matters here: the kernel is bound by + memory throughput rather than per-wave latency, so too few waves leaves the + memory system idle and too many just queue. Accounting for the CTA's wave + count keeps the target fixed across shapes. Returns >= 1. """ tile_q_count = max(1, batch_size * next_n) if total_cu is None: @@ -348,11 +341,10 @@ def launch_fp8_paged_mqa_logits( # Kernel variants: single-token-per-block, wave-split only ("paged_w"). # WPB must divide the column-tile count BKV/16 (=8 at the default BKV=128). # -# WPB=1 (one wave per CTA) is the default: with SplitKV set to hit the same -# resident wave count, splitting a CTA across 2 or 4 waves is 4-14% slower on -# every shape measured -- a 1-wave CTA schedules more freely and pays no -# whole-CTA granularity cost, and the waves have no work to share anyway (each -# owns a disjoint set of KV columns, with no barrier or LDS between them). +# WPB=1 (one wave per CTA) is the default: a 1-wave CTA schedules more freely +# and pays no whole-CTA granularity cost, and the waves have no work to share +# anyway, each owning a disjoint set of KV columns with no barrier or LDS +# between them. KERNEL_VARIANTS = tuple(f"paged_w{w}" for w in (1, 2, 4)) DEFAULT_VARIANT = "paged_w1" @@ -451,8 +443,7 @@ def flydsl_fp8_paged_mqa_logits( max_model_len: out_logits column count. SplitKV: KV-column split count (grid = split_kv*batch*next_n). None => sized for WavePerEU waves per SIMD; 1 disables splitting. - WavePerEU: target resident waves per SIMD (default 2, where the measured - curve bottoms out). + WavePerEU: target resident waves per SIMD for the auto-SplitKV formula. TotalCuCount: CU count for that formula; None => query the device. ChunkK: KV tile width, a multiple of MFMA_N=32. Defaults to 128. diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 365a04dabe6..4ee4e9ea63c 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -187,9 +187,8 @@ def _build_inputs( # Block table: one entry per KVBlockSize-token block, handed out of a # shuffled pool in sequence order. The pool is sized from max_model_len, so - # sequences share blocks once the batch needs more than it holds -- harmless - # here (and extra coverage of block reuse), unlike in the benchmarks where - # it would shrink the measured working set into cache. + # sequences share blocks once the batch needs more than it holds. That is + # deliberate: it costs nothing for correctness and exercises block reuse. blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size max_block_len = int(blocks_per_seq.max().item()) pool = list(range(num_blocks)) From 2b3ad3ffdbd9427cce0023bd61680aa44f7f54e9 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Sat, 15 Aug 2026 07:28:23 +0000 Subject: [PATCH 20/52] docs(flydsl): fix stale --variant default in the paged bench help The help text still advertised paged_w4 after DEFAULT_VARIANT moved to paged_w1. Read it from the constant so it cannot drift again. Co-authored-by: Cursor --- .../op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py b/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py index d6d23dcc08f..6f8cb9bab39 100644 --- a/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py @@ -161,7 +161,10 @@ def main(): parser.add_argument("--wave-per-eu", type=int, default=2) parser.add_argument("--chunk-k", type=int, default=128, help="KV tile width") parser.add_argument( - "--variant", type=str, default=None, help="paged_w tag (default paged_w4)" + "--variant", + type=str, + default=None, + help=f"paged_w tag; None uses the kernel default ({DEFAULT_VARIANT})", ) parser.add_argument("--var-ratio", type=float, default=0.0) parser.add_argument("--seed", type=int, default=0) From 22e4a84d71839a79763421785eea494a826af176 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Sat, 15 Aug 2026 07:42:51 +0000 Subject: [PATCH 21/52] refactor: move the FlyDSL benchmarks out of op_benchmarks/triton They were sitting in the Triton benchmark directory alongside 73 unrelated files, which is misleading -- neither uses Triton. op_benchmarks/flydsl/ now holds both, matching how op_tests/flydsl_tests/ already separates the FlyDSL correctness tests. No __init__.py needed: op_benchmarks subdirectories are namespace packages, which is also why the vs-gluon benchmark keeps importing kv_cache_cast_to_fp8 from op_benchmarks.triton.bench_deepgemm_attention -- that file is genuinely a Triton benchmark and stays put. Paths updated in the bench's own rocprofv3 example and in pr.md. Both run clean from the new location: microbenchmark 27.8us with its correctness gate passing, head-to-head 3.11x at B=16. Co-authored-by: Cursor --- .../{triton => flydsl}/bench_flydsl_paged_mqa_logits.py | 2 +- .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename op_tests/op_benchmarks/{triton => flydsl}/bench_flydsl_paged_mqa_logits.py (99%) rename op_tests/op_benchmarks/{triton => flydsl}/bench_flydsl_vs_gluon_paged_mqa_logits.py (100%) diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py b/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py similarity index 99% rename from op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py rename to op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py index 6f8cb9bab39..9404c7aae52 100644 --- a/op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py @@ -8,7 +8,7 @@ PYTHONPATH=. rocprofv3 --kernel-trace --output-format csv \\ --output-directory .rocprofv3 -- \\ - python op_tests/op_benchmarks/triton/bench_flydsl_paged_mqa_logits.py \\ + python op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py \\ --heads 64 --head-dim 128 --kv-len 32768 --batch 16 --next-n 2 \\ --kv-block-size 64 --preshuffle --iters 200 --no-check diff --git a/op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py similarity index 100% rename from op_tests/op_benchmarks/triton/bench_flydsl_vs_gluon_paged_mqa_logits.py rename to op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py From b11cd5a3e2cbbc5194541204862927e56d9b705d Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Sat, 15 Aug 2026 08:50:39 +0000 Subject: [PATCH 22/52] style(flydsl): satisfy the Black and Ruff pre-checks on PR #4221 CI reported five failures; only two are real. check-signal waits on the pre-checks workflow and Aiter Test Gate depends on check-signal, so the gate failed in four seconds without running a single GPU test. Fixing style unblocks all three. These findings were invisible locally because I had been diffing ruff against the branch's own history. CI runs it through reviewdog with filter-mode=diff_context against upstream/main, where none of these files exist -- so every line is an added line and every finding is in the diff. Seven auto-fixed: import ordering in five files plus __all__ ordering in aiter/ops/flydsl/__init__.py. Five needed a decision, taken to match what the repo already does rather than inventing a convention -- it carries 214 BLE001, 21 B008 and 19 PLR0124 suppressions, including this exact `fx.Stream(None)` default in moe_route_maps.py and this exact NaN-check idiom in test_fused_qk_norm_rope_group_quant.py. So noqa for the blind except, the FlyDSL stream default and the self-comparison, and the two unused `heads` unpackings renamed to `_heads`. Black reformatted one file. Verified as CI runs them, repo-wide: `black --check .` leaves 1063 files unchanged, and `ruff check .` leaves five findings, all in upstream files this PR does not touch. Behaviour unchanged: 35/35 correctness at 3.97e-4, 27.93us, and the public exports still resolve after __init__ was resorted. Co-authored-by: Cursor --- aiter/ops/flydsl/__init__.py | 20 +++++++++++------- .../kernels/mqa_logits/_mqa_logits_common.py | 5 ++--- .../mqa_logits/fp8_paged_mqa_logits.py | 18 +++++++--------- .../test_flydsl_fp8_paged_mqa_logits.py | 7 ++++--- .../flydsl/bench_flydsl_paged_mqa_logits.py | 21 +++++++++++-------- .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 12 ++++++----- 6 files changed, 45 insertions(+), 38 deletions(-) diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 5953c93cc3b..92a202e1bb8 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -48,6 +48,15 @@ from .kernels.mqa_logits.fp8_mqa_logits import ( flydsl_fp8_mqa_logits, ) + from .kernels.mqa_logits.fp8_paged_mqa_logits import ( + DEFAULT_VARIANT as FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT, + ) + from .kernels.mqa_logits.fp8_paged_mqa_logits import ( + KERNEL_VARIANTS as FP8_PAGED_MQA_LOGITS_VARIANTS, + ) + from .kernels.mqa_logits.fp8_paged_mqa_logits import ( + flydsl_fp8_paged_mqa_logits, + ) from .kernels.mqa_logits.pa_mqa_logits_fp4 import ( flydsl_pa_mqa_logits_fp4, ) @@ -56,11 +65,6 @@ flydsl_pa_mqa_logits_fp4_prefill, flydsl_pa_mqa_logits_fp4_varqlen, ) - from .kernels.mqa_logits.fp8_paged_mqa_logits import ( - flydsl_fp8_paged_mqa_logits, - KERNEL_VARIANTS as FP8_PAGED_MQA_LOGITS_VARIANTS, - DEFAULT_VARIANT as FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT, - ) from .kernels.qk_norm_rope_quant import flydsl_qk_norm_rope_quant from .mla_reduce_kernels import flydsl_mla_reduce_v1 from .moe_kernels import flydsl_moe_stage1, flydsl_moe_stage2 @@ -70,9 +74,12 @@ __all__ += [ "FP8_MQA_LOGITS_DEFAULT_VARIANT", "FP8_MQA_LOGITS_VARIANTS", + "FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT", + "FP8_PAGED_MQA_LOGITS_VARIANTS", "compute_varqlen_windows", "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", + "flydsl_fp8_paged_mqa_logits", "flydsl_hgemm", "flydsl_mla_reduce_v1", "flydsl_moe_stage1", @@ -82,8 +89,5 @@ "flydsl_pa_mqa_logits_fp4_varqlen", "flydsl_preshuffle_gemm_a8", "flydsl_qk_norm_rope_quant", - "flydsl_fp8_paged_mqa_logits", - "FP8_PAGED_MQA_LOGITS_VARIANTS", - "FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT", # "flydsl_gdr_decode", ] diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index 46f4a8f233e..d5ca9a9e84b 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -16,9 +16,8 @@ from functools import lru_cache -import torch - import flydsl.expr as fx +import torch from flydsl.expr import arith, range_constexpr, rocdl from flydsl.expr.numeric import ArithValue from flydsl.expr.typing import T @@ -72,7 +71,7 @@ def device_cu_count(device_index: int) -> int: """Compute-unit count for a CUDA/HIP device (cached); 304 if unavailable.""" try: return torch.cuda.get_device_properties(device_index).multi_processor_count - except Exception: + except Exception: # noqa: BLE001 return 304 diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index f4bf9121345..5da73c497cb 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -35,23 +35,22 @@ import os from functools import lru_cache -import torch - -from aiter.jit.utils.chip_info import get_gfx - import flydsl.compiler as flyc import flydsl.expr as fx +import torch +from flydsl._mlir import ir +from flydsl._mlir.dialects import scf from flydsl.expr import arith, range_constexpr from flydsl.expr.typing import T -from flydsl._mlir.dialects import scf -from flydsl._mlir import ir + +from aiter.jit.utils.chip_info import get_gfx from ..tensor_shim import GTensor, _run_compiled, _to_raw from ._mqa_logits_common import ( + DEFAULT_COMPILE_HINTS, + MFMA_K, MFMA_M, MFMA_N, - MFMA_K, - DEFAULT_COMPILE_HINTS, Vec, device_cu_count, load_pack_kv, @@ -64,7 +63,6 @@ umod, ) - _SIMDS_PER_CU = 4 @@ -313,7 +311,7 @@ def launch_fp8_paged_mqa_logits( index_dim: fx.Int32, max_block_len: fx.Int32, stride_out: fx.Int32, - stream: fx.Stream = fx.Stream(None), + stream: fx.Stream = fx.Stream(None), # noqa: B008 ): gx = arith.index_cast(T.index, _to_raw(grid_blocks)) kernel._func.__name__ = _kname diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 4ee4e9ea63c..c5dd36df9c3 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -13,13 +13,14 @@ import random from typing import NamedTuple -import aiter import pandas as pd import torch + +import aiter from aiter import dtypes -from aiter.test_common import benchmark, checkAllclose from aiter.jit.utils.chip_info import get_gfx from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from aiter.test_common import benchmark, checkAllclose torch.set_default_device("cuda") @@ -106,7 +107,7 @@ def ref_fp8_paged_mqa_logits( masked by ``p <= context_len - next_n + n``. Position ``p`` resolves to ``(block_tables[b, p//KVB], p%KVB)`` in the block-flat co-packed cache. """ - batch_size, next_n, heads, dim = q.size() + batch_size, next_n, _heads, dim = q.size() num_blocks = kv_cache_fp8.shape[0] index_dim = kv_cache_fp8.shape[-1] # Reconstruct the block-flat block: [KVB keys (dim bytes each)][KVB f32 scales]. diff --git a/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py b/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py index 9404c7aae52..311516a2ee9 100644 --- a/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py @@ -30,19 +30,19 @@ import torch -from aiter.test_common import run_perftest from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits from aiter.ops.flydsl.kernels.mqa_logits.fp8_paged_mqa_logits import DEFAULT_VARIANT +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from aiter.test_common import run_perftest # Reuse the correctness test's layout/oracle helpers so the profiled kernel sees # exactly the tested co-pack + preshuffle byte layout. from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( + calc_diff, kv_cache_cast_to_fp8, - ref_fp8_paged_mqa_logits, preshuffle_kv_data, - calc_diff, + ref_fp8_paged_mqa_logits, ) torch.set_default_device("cuda") @@ -152,11 +152,12 @@ def main(): action="store_true", help="use the shuffle_weight(16x16) KV layout (needs kv-block-size %% 16 == 0)", ) + parser.add_argument("--q-dtype", type=str, default="fnuz", choices=["fnuz", "fn"]) parser.add_argument( - "--q-dtype", type=str, default="fnuz", choices=["fnuz", "fn"] - ) - parser.add_argument( - "--split-kv", type=int, default=0, help="0 == auto (host formula); else override" + "--split-kv", + type=int, + default=0, + help="0 == auto (host formula); else override", ) parser.add_argument("--wave-per-eu", type=int, default=2) parser.add_argument("--chunk-k", type=int, default=128, help="KV tile width") @@ -183,7 +184,9 @@ def main(): "--num-iters", type=int, default=101, help="run_perftest iters for --time" ) parser.add_argument( - "--time", action="store_true", help="also report device self-time (run_perftest)" + "--time", + action="store_true", + help="also report device self-time (run_perftest)", ) parser.add_argument( "--no-check", action="store_true", help="skip the one-shot correctness gate" diff --git a/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py index 3a5dd9959be..2e381b20daa 100644 --- a/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py +++ b/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py @@ -29,11 +29,11 @@ import torch -from aiter.test_common import run_perftest from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits +from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from aiter.test_common import run_perftest # same fp8 co-pack builder the production paged benchmark uses from op_tests.op_benchmarks.triton.bench_deepgemm_attention import ( @@ -57,7 +57,7 @@ def ref_fp8_paged_mqa_logits( q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype ): """Vectorized torch reference (dequantizes the co-packed fp8 cache).""" - batch_size, next_n, heads, dim = q.size() + batch_size, next_n, _heads, dim = q.size() kvv, scale = kv_cache_fp8[..., :dim], kv_cache_fp8[..., dim:] scale = scale.contiguous().view(torch.float) qf = q.float() @@ -328,7 +328,9 @@ def fn_fly(): "fly_min_ms": fly_min / 1e3, "fly_max_ms": fly_max / 1e3, "fly/ref": ( - (fly_med / ref_med) if ref_med == ref_med and ref_med > 0 else float("nan") + (fly_med / ref_med) + if ref_med == ref_med and ref_med > 0 # noqa: PLR0124 (NaN check) + else float("nan") ), "ref_diff": ref_diff, "fly_diff": fly_diff, From b93e68c0e24bd121dfa1d162ad709da38b412833 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Sat, 15 Aug 2026 13:35:48 +0000 Subject: [PATCH 23/52] refactor(flydsl): drop dead conversion plumbing from the paged kernel convert_q_fn and convert_kv_fn were threaded through the kernel builder, the compile cache and the host, and neither is ever read inside the emitted kernel -- they only appended a name tag. convert_kv_fn was additionally hardcoded False at every call site. Only scale_mul reaches codegen, so it is now the sole carrier of the FN->FNUZ compensation and tags the kernel name itself. That exposed a docstring that was simply untrue: compile_fp8_paged_mqa_logits claimed these flags "mark an FP8 FN operand whose -0 (0x80) byte the kernel patches to FNUZ +0". No such patch exists on this path -- the vector<8xi32> operand has no byte-patch, which the kernel body's own comment admitted. The host comment now states the actual gfx942 behaviour, including that FN -0 decodes to FNUZ NaN unpatched. Flagged rather than fixed: it needs gfx942 to validate. Also: mfma_head_reduce now derives res_ty and f32_0 and defaults fm_fast, so the call site is one line instead of eleven and the kernel drops two locals; make_out_row_view takes the i32 stride and widens it itself. Both are safe because _mqa_logits_common turns out to have exactly one consumer -- upstream's dense fp8_mqa_logits carries its own copies and does not import it, so the module docstring claiming otherwise is corrected too. 569 -> 535 lines in the kernel. ISA byte-identical to HEAD (669 instructions), 35/35 correctness at 3.97e-4, 27.92us, black and ruff clean. Co-authored-by: Cursor --- .../kernels/mqa_logits/_mqa_logits_common.py | 20 +++---- .../mqa_logits/fp8_paged_mqa_logits.py | 59 ++++--------------- 2 files changed, 23 insertions(+), 56 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index d5ca9a9e84b..b5f4e10b751 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -1,14 +1,13 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Shared FlyDSL building blocks for the dense and paged FP8 MQA-logits kernels. - -Both compute the same core per query row / KV column:: +"""FlyDSL building blocks for the paged FP8 MQA-logits kernel:: logits[row, n] = sum_h ReLU( * kv_scale[n]) * weights[row, h] -and differ only in how K / kv_scale are addressed (contiguous vs paged gather) -and in the windowing / output layout. Everything else lives here. +Split out of ``fp8_paged_mqa_logits`` to keep the gather layouts, the MFMA +head-reduce and the output view separable from the kernel body. The dense +``fp8_mqa_logits`` carries its own copies and does not import this. """ # No `from __future__ import annotations`: FlyDSL arg typing needs real @@ -141,11 +140,12 @@ def _load_i64(sub): ).bitcast(fx.Int32) -def make_out_row_view(logits, stride_i64, row_i32): +def make_out_row_view(logits, stride_out, row_i32): """1-D output GTensor for ``row_i32``, with the row byte base folded into the pointer in i64 so the column offset stays i32. A 2-D (row, col) view would compute ``row*stride + col`` in i32 and silently overflow past 2^31. """ + stride_i64 = arith.extui(T.i64, _to_raw(stride_out)) _ri64 = arith.extui(T.i64, _to_raw(row_i32)) _byte = arith.muli(arith.muli(_ri64, stride_i64), arith.constant(4, type=T.i64)) _idx = arith.index_cast(T.index, _byte) @@ -160,11 +160,9 @@ def mfma_head_reduce( *, m_tiles, k_steps, - res_ty, - f32_0, - fm_fast, - mfma_fn=MFMA_FN, dreg_count=16, + fm_fast=arith.FastMathFlags.fast, + mfma_fn=MFMA_FN, ): """One column's logit: fp8 MFMA over heads, ReLU * weight sum, kv-scale, in-wave head reduce. Returns the reduced f32 ``col_sum`` (raw ir value). @@ -173,6 +171,8 @@ def mfma_head_reduce( ReLU is positively homogeneous, so ReLU(s*x) = s*ReLU(x) and hoisting it out of the head sum is exact. """ + res_ty = Vec.make_type(dreg_count, fx.Float32) + f32_0 = arith.constant(0.0, type=T.f32) col_sum = _to_raw(f32_0) for mi in range_constexpr(m_tiles): acc = Vec.filled(dreg_count, 0.0, fx.Float32) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 5da73c497cb..81ecc9cf49f 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -51,7 +51,6 @@ MFMA_K, MFMA_M, MFMA_N, - Vec, device_cu_count, load_pack_kv, load_pack_v8i32, @@ -99,8 +98,6 @@ def _build_paged_kernel( kv_block_size: int = 1, preshuffle: bool = False, scale_mul: float = 1.0, - convert_q_fn: bool = False, - convert_kv_fn: bool = False, ): """Paged MQA-logits kernel with SplitKV column parallelism. @@ -135,16 +132,9 @@ def _build_paged_kernel( N_TILES_PER_WAVE = N_TILES // WPB fm_fast = arith.FastMathFlags.fast - - _cvt_tag = "" - if convert_q_fn: - _cvt_tag += "_cq" - if convert_kv_fn: - _cvt_tag += "_ck" - _ps_tag = "_ps" if preshuffle else "" _kname = ( f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_kvb{KVB}_w{WPB}" - f"{_ps_tag}{_cvt_tag}_flydsl" + f"{'_ps' if preshuffle else ''}{'' if scale_mul == 1.0 else '_s2'}_flydsl" ) @flyc.kernel(name=_kname, known_block_size=[MR_BLOCK_THREADS, 1, 1]) @@ -165,8 +155,6 @@ def kernel( max_block_len: fx.Int32, # kv_indices row width stride_out: fx.Int32, # out_logits.stride(0) == max_model_len ): - f32_0 = arith.constant(0.0, type=T.f32) - mfma_res_ty = Vec.make_type(16, fx.Float32) scale_mul_c = arith.constant(float(scale_mul), type=T.f32) tid = fx.thread_idx.x @@ -201,13 +189,11 @@ def kernel( q_limit = context_length - next_n + pid_next_n out_row = pid_batch * next_n + pid_next_n - stride_out_i64 = arith.extui(T.i64, _to_raw(stride_out)) - out_row_t = make_out_row_view(out_logits, stride_out_i64, out_row) + out_row_t = make_out_row_view(out_logits, stride_out, out_row) # ---- Preload Q frags + weights for the single query row ---- # The A-operand layout is per in-wave lane, so `lane` indexes Q, not - # `tid`. convert_q_fn is always False here: the v8i32 operand has no - # fn_to_fnuz path, and gfx950 is native fp8 anyway. + # `tid`. q_row_base = pid_batch * stride_q_batch + pid_next_n * stride_q_next_n a_pack = [[None] * K_STEPS for _ in range_constexpr(M_TILES)] for mi in range_constexpr(M_TILES): @@ -274,16 +260,7 @@ def kernel( ).result col_sum = mfma_head_reduce( - a_pack, - b_col, - w_frag, - kv_scale, - m_tiles=M_TILES, - k_steps=K_STEPS, - res_ty=mfma_res_ty, - f32_0=f32_0, - fm_fast=fm_fast, - dreg_count=_DREG, + a_pack, b_col, w_frag, kv_scale, m_tiles=M_TILES, k_steps=K_STEPS ) # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. @@ -377,16 +354,13 @@ def compile_fp8_paged_mqa_logits( preshuffle: bool = False, variant: str = DEFAULT_VARIANT, scale_mul: float = 1.0, - convert_q_fn: bool = False, - convert_kv_fn: bool = False, ): """Return a cached, compiled FlyDSL paged launcher for the given config. ``num_heads``/``head_size``/``kv_block_size``/``preshuffle`` are compile-time - constants; ``variant`` is a ``paged_w`` tag; ``convert_q_fn``/ - ``convert_kv_fn`` mark an FP8 FN operand whose -0 (0x80) byte the kernel - patches to FNUZ +0, and ``scale_mul`` folds the 2x-per-converted-operand - compensation into the co-packed scale. + constants and ``variant`` is a ``paged_w`` tag. ``scale_mul`` folds a + constant into the co-packed dequant scale; it carries the gfx942 FN->FNUZ + 2x compensation. """ launcher = _build_paged_kernel( num_heads=num_heads, @@ -396,8 +370,6 @@ def compile_fp8_paged_mqa_logits( kv_block_size=kv_block_size, preshuffle=preshuffle, scale_mul=scale_mul, - convert_q_fn=convert_q_fn, - convert_kv_fn=convert_kv_fn, ) launcher.compile_hints = dict(DEFAULT_COMPILE_HINTS) return launcher @@ -498,15 +470,12 @@ def flydsl_fp8_paged_mqa_logits( _fn, ), f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" - # gfx942 fp8 MFMA reads operands as e4m3 FNUZ (bias 8). Q may arrive as FN - # (OCP, bias 7); patch its 0x80 byte and undo the implied 2x via scale_mul. - # The co-packed KV in the cache is arch-native fp8 (fnuz on gfx942, fn on - # gfx950), so it is never converted. - gfx = get_gfx() - convert_q_fn = gfx == "gfx942" and q_fp8.dtype == _fn - convert_kv_fn = False - scale_mul = 2.0 if convert_q_fn else 1.0 - + # gfx942's fp8 MFMA reads operands as e4m3 FNUZ (bias 8), so an FN (OCP, + # bias 7) Q byte decodes to half its value there; scale_mul undoes that on + # the co-packed scale. NOTE: FN -0 (0x80) is FNUZ NaN and is *not* patched + # -- the vector<8xi32> operand has no byte-patch path. The KV cache is + # always arch-native fp8, so it never needs either. + scale_mul = 2.0 if (get_gfx() == "gfx942" and q_fp8.dtype == _fn) else 1.0 variant = _resolve_variant(variant) launcher = compile_fp8_paged_mqa_logits( @@ -517,8 +486,6 @@ def flydsl_fp8_paged_mqa_logits( preshuffle=bool(Preshuffle), variant=variant, scale_mul=scale_mul, - convert_q_fn=convert_q_fn, - convert_kv_fn=convert_kv_fn, ) # Co-packed cache -> raw uint8 byte view [num_blocks * KVBlockSize * index_dim]. From 52afe43d1c58b16bb7b447f9c69a164eb10d1180 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Wed, 19 Aug 2026 09:37:06 +0000 Subject: [PATCH 24/52] fix(flydsl): pass buffer load/store aux as an i32 SSA operand upstream/main treated aux as an optional IntegerAttr, but RawPtrBufferLoadOp expects it as the fourth operand Value. With cache_modifier=0 that left aux None and broke JIT compilation of kernels using GTensor buffer loads. Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/buffer_ops.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/aiter/ops/flydsl/kernels/buffer_ops.py b/aiter/ops/flydsl/kernels/buffer_ops.py index 0399b84f690..cb04c916fce 100644 --- a/aiter/ops/flydsl/kernels/buffer_ops.py +++ b/aiter/ops/flydsl/kernels/buffer_ops.py @@ -593,11 +593,8 @@ def buffer_load( soffset = _create_i32_constant(soffset_bytes) else: soffset = _to_i32_offset(_unwrap_value(soffset_bytes)) - aux_attr = ( - ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) - if cache_modifier - else None - ) + # aux is the 4th SSA operand (i32), not an op attribute — pass 0 when unused. + aux = _create_i32_constant(cache_modifier) # Emit buffer load load_op = rocdl.RawPtrBufferLoadOp( @@ -605,7 +602,7 @@ def buffer_load( rsrc, offset, soffset, - aux=aux_attr, + aux, ) return load_op.result @@ -683,11 +680,8 @@ def buffer_store( soffset = _create_i32_constant(int(soffset_bytes)) else: soffset = _to_i32_offset(_unwrap_value(soffset_bytes)) - aux_attr = ( - ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) - if cache_modifier - else None - ) + # aux is the 4th SSA operand (i32), not an op attribute — pass 0 when unused. + aux = _create_i32_constant(cache_modifier) # Emit buffer store rocdl.RawPtrBufferStoreOp( @@ -695,5 +689,5 @@ def buffer_store( rsrc, offset, soffset, - aux=aux_attr, + aux, ) From 9a9c7f9684fd9604e0e0800588c7ace9e26c2fc0 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Wed, 19 Aug 2026 11:08:00 +0000 Subject: [PATCH 25/52] fix(flydsl): align buffer_ops and MFMA helpers with flydsl 0.3.1 Revert buffer load/store aux to the upstream IntegerAttr form expected by flydsl 0.3.1, and route scaled MFMA through the expr wrapper so cbsz/blgp attributes lower correctly during paged MQA kernel JIT. Co-authored-by: Cursor --- aiter/ops/flydsl/kernels/buffer_ops.py | 18 ++++++++++++------ .../kernels/mqa_logits/_mqa_logits_common.py | 16 +++++++--------- 2 files changed, 19 insertions(+), 15 deletions(-) diff --git a/aiter/ops/flydsl/kernels/buffer_ops.py b/aiter/ops/flydsl/kernels/buffer_ops.py index cb04c916fce..0399b84f690 100644 --- a/aiter/ops/flydsl/kernels/buffer_ops.py +++ b/aiter/ops/flydsl/kernels/buffer_ops.py @@ -593,8 +593,11 @@ def buffer_load( soffset = _create_i32_constant(soffset_bytes) else: soffset = _to_i32_offset(_unwrap_value(soffset_bytes)) - # aux is the 4th SSA operand (i32), not an op attribute — pass 0 when unused. - aux = _create_i32_constant(cache_modifier) + aux_attr = ( + ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) + if cache_modifier + else None + ) # Emit buffer load load_op = rocdl.RawPtrBufferLoadOp( @@ -602,7 +605,7 @@ def buffer_load( rsrc, offset, soffset, - aux, + aux=aux_attr, ) return load_op.result @@ -680,8 +683,11 @@ def buffer_store( soffset = _create_i32_constant(int(soffset_bytes)) else: soffset = _to_i32_offset(_unwrap_value(soffset_bytes)) - # aux is the 4th SSA operand (i32), not an op attribute — pass 0 when unused. - aux = _create_i32_constant(cache_modifier) + aux_attr = ( + ir.IntegerAttr.get(ir.IntegerType.get_signless(32), cache_modifier) + if cache_modifier + else None + ) # Emit buffer store rocdl.RawPtrBufferStoreOp( @@ -689,5 +695,5 @@ def buffer_store( rsrc, offset, soffset, - aux, + aux=aux_attr, ) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index b5f4e10b751..8f7aae93990 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -33,17 +33,15 @@ def _mfma_scale_f32_32x32x64_fp8(result_type, operands): - """operands: [a, b, c, cbsz, blgp, opsel_a, scale_a, opsel_b, scale_b].""" - from flydsl.expr.rocdl import _unwrap_mfma_operand - - a = _unwrap_mfma_operand(operands[0]) - b = _unwrap_mfma_operand(operands[1]) - c = _unwrap_mfma_operand(operands[2]) + """operands: [a, b, c] or [a, b, c, cbsz, blgp, opsel_a, scale_a, opsel_b, scale_b].""" # Neutral e8m0 scale: 4 packed bytes of 127 => 2^(127-127) = 1.0 each. neutral = arith.constant(0x7F7F7F7F, type=T.i32) - return rocdl.mfma_scale_f32_32x32x64_f8f6f4_( - result_type, a, b, c, 0, 0, 0, neutral, 0, neutral - ) + ops = list(operands) + if len(ops) < 9: + while len(ops) < 6: + ops.append(0) + ops.extend([neutral, 0, neutral]) + return rocdl.mfma_scale_f32_32x32x64_f8f6f4(result_type, ops) MFMA_FN = _mfma_scale_f32_32x32x64_fp8 From acf027710da3898a93780c2499ad7e90d00818df Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Thu, 20 Aug 2026 09:39:52 +0000 Subject: [PATCH 26/52] perf(flydsl): software-pipeline block-table and K0 prefetch in paged MQA logits Prefetch kv_indices one tile ahead (skip reload when ind_off unchanged) and carry K-step-0 across the inner column loop to hide VMEM latency on the production kvb64/preshuffle decode path. Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits.py | 79 ++++++++++++++----- 1 file changed, 59 insertions(+), 20 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 81ecc9cf49f..7d12e5afb09 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -236,40 +236,79 @@ def kernel( ): col0 = fx.Int32(arith.index_cast(T.i32, iv)) wave_ni_base = wave * N_TILES_PER_WAVE - for ni in range_constexpr(N_TILES_PER_WAVE): - abs_ni = wave_ni_base + ni - col = col0 + abs_ni * MFMA_N + lane_mod_N - # Clamp the gather index to a valid token (mask handles the rest). - col_c = fx.Int32(arith.minsi(_to_raw(col), _to_raw(ctx_m1))) - # KVB is a compile-time constant, so at KVB==1 the div/mod/mul - # fold away on their own. Deliberately no Python `if` for that - # case: the FlyDSL AST rewriter would turn it into a runtime - # scf.if and drop the branch-local addresses. - tok_in_block = umod(col_c, KVB) - ind_off = pid_batch * max_block_len + udiv(col_c, KVB) - physical = fx.Int32(ind_t[ind_off]) - - b_col = [ - load_pack_kv(kv_key, physical, tok_in_block, kk, lane_div_N) - for kk in range_constexpr(K_STEPS) + + def _col_for(abs_ni): + return col0 + abs_ni * MFMA_N + lane_mod_N + + def _col_c_for(abs_ni): + return fx.Int32( + arith.minsi(_to_raw(_col_for(abs_ni)), _to_raw(ctx_m1)) + ) + + def _ind_off_for(abs_ni): + return pid_batch * max_block_len + udiv(_col_c_for(abs_ni), KVB) + + def _tok_in_block_for(abs_ni): + return umod(_col_c_for(abs_ni), KVB) + + def _load_physical(ind_off): + return fx.Int32(ind_t[ind_off]) + + def _prefetch_physical(physical_cur, ind_off_cur, ind_off_next): + same = _to_raw(ind_off_next) == _to_raw(ind_off_cur) + phys_raw = _to_raw(physical_cur) + if_op = scf.IfOp(same, results_=[T.i32], has_else=True) + with ir.InsertionPoint(if_op.then_block): + scf.YieldOp([phys_raw]) + with ir.InsertionPoint(if_op.else_block): + scf.YieldOp([_to_raw(_load_physical(ind_off_next))]) + return fx.Int32(if_op.results[0]) + + def _b_col_for_tile(physical, tok_in_block, b0_carried): + return [ + b0_carried, + load_pack_kv(kv_key, physical, tok_in_block, 1, lane_div_N), ] + + def _process_tile_with_b0(abs_ni, physical, b0_carried): + col = _col_for(abs_ni) + tok_in_block = _tok_in_block_for(abs_ni) + b_col = _b_col_for_tile(physical, tok_in_block, b0_carried) kv_scale = _to_raw(fx.Float32(kv_scale_v[physical, tok_in_block])) if scale_mul != 1.0: kv_scale = arith.MulFOp( kv_scale, scale_mul_c, fastmath=fm_fast ).result - col_sum = mfma_head_reduce( a_pack, b_col, w_frag, kv_scale, m_tiles=M_TILES, k_steps=K_STEPS ) - - # Only lane_div_N==0 lanes hold the MFMA_N distinct columns. - # Masked positions keep the caller's -inf prefill. is_writer = _to_raw((lane_div_N == 0) & (col <= q_limit)) with ir.InsertionPoint(scf.IfOp(is_writer).then_block): out_row_t[col] = fx.Float32(col_sum) scf.YieldOp([]) + def _prefetch_k0(abs_ni, physical): + return load_pack_kv( + kv_key, physical, _tok_in_block_for(abs_ni), 0, lane_div_N + ) + + physical_carry = _load_physical(_ind_off_for(wave_ni_base + 0)) + b0_carry = _prefetch_k0(wave_ni_base + 0, physical_carry) + for ni in range_constexpr(N_TILES_PER_WAVE): + abs_ni = wave_ni_base + ni + physical = physical_carry + if ni + 1 < N_TILES_PER_WAVE: + abs_ni_next = wave_ni_base + ni + 1 + physical_carry = _prefetch_physical( + physical, + _ind_off_for(abs_ni), + _ind_off_for(abs_ni_next), + ) + _process_tile_with_b0(abs_ni, physical, b0_carry) + if ni + 1 < N_TILES_PER_WAVE: + abs_ni_next = wave_ni_base + ni + 1 + b0_carry = _prefetch_k0(abs_ni_next, physical_carry) + @flyc.jit def launch_fp8_paged_mqa_logits( Q: fx.Tensor, From 99b0ec3d7287a18c59b44785c15fd9398d9869e0 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Thu, 20 Aug 2026 09:45:54 +0000 Subject: [PATCH 27/52] style(flydsl): fix Black/Ruff on paged MQA logits prefetch helpers Bind loop SSA values via default args (B023) and apply Black formatting. Co-authored-by: Cursor --- .../flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 7d12e5afb09..2a28b06bc32 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -237,12 +237,13 @@ def kernel( col0 = fx.Int32(arith.index_cast(T.i32, iv)) wave_ni_base = wave * N_TILES_PER_WAVE - def _col_for(abs_ni): - return col0 + abs_ni * MFMA_N + lane_mod_N + # Helpers close over this chunk's SSA values; default args silence B023. + def _col_for(abs_ni, _col0=col0): + return _col0 + abs_ni * MFMA_N + lane_mod_N - def _col_c_for(abs_ni): + def _col_c_for(abs_ni, _ctx_m1=ctx_m1): return fx.Int32( - arith.minsi(_to_raw(_col_for(abs_ni)), _to_raw(ctx_m1)) + arith.minsi(_to_raw(_col_for(abs_ni)), _to_raw(_ctx_m1)) ) def _ind_off_for(abs_ni): From 0725f0d40d189dbad73178b058e9a6eed142b07a Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 25 Aug 2026 11:54:07 +0000 Subject: [PATCH 28/52] address comments by Sami --- .../kernels/mqa_logits/_mqa_logits_common.py | 41 +- .../mqa_logits/fp8_paged_mqa_logits.py | 95 ++-- .../test_flydsl_fp8_paged_mqa_logits.py | 473 +++++++++++++----- .../flydsl/bench_flydsl_paged_mqa_logits.py | 335 ------------- .../bench_flydsl_vs_gluon_paged_mqa_logits.py | 427 ---------------- 5 files changed, 412 insertions(+), 959 deletions(-) delete mode 100644 op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py delete mode 100644 op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index 8f7aae93990..d6577443e18 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -1,32 +1,24 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""FlyDSL building blocks for the paged FP8 MQA-logits kernel:: +"""FlyDSL building blocks for the paged FP8 MQA-logits kernel (gfx950 / CDNA4). logits[row, n] = sum_h ReLU( * kv_scale[n]) * weights[row, h] - -Split out of ``fp8_paged_mqa_logits`` to keep the gather layouts, the MFMA -head-reduce and the output view separable from the kernel body. The dense -``fp8_mqa_logits`` carries its own copies and does not import this. """ -# No `from __future__ import annotations`: FlyDSL arg typing needs real -# annotation objects, not PEP 563 strings. - from functools import lru_cache import flydsl.expr as fx import torch from flydsl.expr import arith, range_constexpr, rocdl -from flydsl.expr.numeric import ArithValue from flydsl.expr.typing import T from ..tensor_shim import GTensor, _to_raw Vec = fx.Vector -# 32x32x64 scaled fp8 MFMA (CDNA4). A/B are vector<8xi32> per lane (32 fp8 -# bytes), C/D are 16 f32 D-regs. lane_div_N = lane//32, lane_mod_N = lane%32. +# 32x32x64 scaled fp8 MFMA (gfx950 / CDNA4). A/B are vector<8xi32> per lane +# (32 fp8 bytes), C/D are 16 f32 D-regs. lane_div_N = lane//32, lane_mod_N = lane%32. MFMA_M = 32 MFMA_N = 32 MFMA_K = 64 @@ -65,18 +57,18 @@ def umod(a, b): @lru_cache(maxsize=8) def device_cu_count(device_index: int) -> int: - """Compute-unit count for a CUDA/HIP device (cached); 304 if unavailable.""" + """Compute-unit count for a CUDA/HIP device (cached); 256 if unavailable.""" try: return torch.cuda.get_device_properties(device_index).multi_processor_count except Exception: # noqa: BLE001 - return 304 + return 256 def load_pack_v8i32(i32_view, byte_off_i32, lane8): """Load 32 fp8 bytes as the vector<8xi32> A/B operand of the 32x32x64 MFMA. Each lane holds 4 K-groups of 8 bytes at ``lane8 + kk*16``. Reads go out as - 2-dword loads because a ``v8i8`` buffer_load fails to lower on gfx942. + 2-dword loads because a ``v8i8`` buffer_load fails to lower. """ def _load_i64(off): @@ -159,34 +151,31 @@ def mfma_head_reduce( m_tiles, k_steps, dreg_count=16, - fm_fast=arith.FastMathFlags.fast, mfma_fn=MFMA_FN, ): """One column's logit: fp8 MFMA over heads, ReLU * weight sum, kv-scale, - in-wave head reduce. Returns the reduced f32 ``col_sum`` (raw ir value). + in-wave head reduce. ``kv_scale`` (>= 0) is applied once to the column sum rather than per head: ReLU is positively homogeneous, so ReLU(s*x) = s*ReLU(x) and hoisting it out of the head sum is exact. """ res_ty = Vec.make_type(dreg_count, fx.Float32) - f32_0 = arith.constant(0.0, type=T.f32) - col_sum = _to_raw(f32_0) + f32_0 = fx.Float32(0.0) + col_sum = f32_0 for mi in range_constexpr(m_tiles): acc = Vec.filled(dreg_count, 0.0, fx.Float32) for kk in range_constexpr(k_steps): acc = mfma_fn(res_ty, [a_row[mi][kk], b_col[kk], acc, 0, 0, 0]) for ii in range_constexpr(dreg_count): - score = Vec(acc)[ii].ir_value() - relu = arith.maximumf(score, _to_raw(f32_0)) - wsc = arith.MulFOp(relu, w_row[mi][ii], fastmath=fm_fast).result - col_sum = arith.AddFOp(col_sum, wsc, fastmath=fm_fast).result - col_sum = arith.MulFOp(col_sum, kv_scale, fastmath=fm_fast).result + score = fx.Float32(Vec(acc)[ii]) + relu = score.maximumf(f32_0) + col_sum = col_sum + relu * w_row[mi][ii] + col_sum = col_sum * fx.Float32(kv_scale) # For 32x32 tile: lane_div_N = lane // 32 (0 or 1); shuffle_xor(32) sums them. - # For 16x16 tile: lane_div_N = lane // 16 (0..3); shuffle_xor(16) + xor(32). shuffles = (32,) if dreg_count == 16 else (16, 32) for sh in shuffles: - peer = _to_raw(ArithValue(col_sum).shuffle_xor(sh, 64)) - col_sum = arith.AddFOp(col_sum, peer, fastmath=fm_fast).result + peer = col_sum.shuffle_xor(sh, 64) + col_sum = col_sum + peer return col_sum diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 2a28b06bc32..242a164546d 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Paged FP8 MQA logits (DeepSeek lightning indexer, decode) -- FlyDSL gfx942/gfx950. +"""Paged FP8 MQA logits (DeepSeek lightning indexer, decode) -- FlyDSL gfx950. Compute for each decode batch element ``b`` and next-n query slot ``n`` (query row ``b*next_n + n``) and KV logical position ``p``:: @@ -38,8 +38,6 @@ import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl._mlir import ir -from flydsl._mlir.dialects import scf from flydsl.expr import arith, range_constexpr from flydsl.expr.typing import T @@ -89,6 +87,18 @@ def _auto_split_kv( _BLOCK_KV = 128 +def _uceildiv(a, b): + a = fx.Int32(a) + b = fx.Int32(b) + return fx.Int32((fx.Uint32(a) + fx.Uint32(b) - 1) // fx.Uint32(b)) + + +def _imin(a, b): + a = fx.Int32(a) + b = fx.Int32(b) + return (a <= b).select(a, b) + + def _build_paged_kernel( *, num_heads: int, @@ -97,7 +107,6 @@ def _build_paged_kernel( waves_per_block: int, kv_block_size: int = 1, preshuffle: bool = False, - scale_mul: float = 1.0, ): """Paged MQA-logits kernel with SplitKV column parallelism. @@ -108,8 +117,6 @@ def _build_paged_kernel( ``preshuffle`` only changes the intra-block key-byte offset -- the block-table gather, scale gather, MFMA and mask are identical either way. - ``scale_mul`` folds the FN->FNUZ 2x compensation into the per-token scale at - compile time, which is exact by ReLU positive-homogeneity. """ H = num_heads D = head_size @@ -131,10 +138,9 @@ def _build_paged_kernel( K_STEPS = D // MFMA_K N_TILES_PER_WAVE = N_TILES // WPB - fm_fast = arith.FastMathFlags.fast _kname = ( f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_kvb{KVB}_w{WPB}" - f"{'_ps' if preshuffle else ''}{'' if scale_mul == 1.0 else '_s2'}_flydsl" + f"{'_ps' if preshuffle else ''}_flydsl" ) @flyc.kernel(name=_kname, known_block_size=[MR_BLOCK_THREADS, 1, 1]) @@ -155,8 +161,6 @@ def kernel( max_block_len: fx.Int32, # kv_indices row width stride_out: fx.Int32, # out_logits.stride(0) == max_model_len ): - scale_mul_c = arith.constant(float(scale_mul), type=T.f32) - tid = fx.thread_idx.x bid = fx.block_idx.x @@ -213,28 +217,21 @@ def kernel( for ii in range_constexpr(_DREG): h_off = (ii % 4) + lane_div_N * 4 + (ii // 4) * 8 h_w = mi * MFMA_M + h_off - w_frag[mi][ii] = _to_raw(fx.Float32(w_t[out_row, h_w])) + w_frag[mi][ii] = fx.Float32(w_t[out_row, h_w]) # ---- SplitKV: each split owns ceil(total_tiles/split_kv) contiguous # BKV-aligned tiles, so the slices are gap-free and disjoint and every # logits column has exactly one writer -- no cross-CTA reduction. A # split past ctx runs zero iterations. ---- - context_chunk_num = arith.ceildivui( - _to_raw(context_length), _to_raw(fx.Int32(BKV)) - ) - split_chunk_num = arith.ceildivui(context_chunk_num, _to_raw(split_kv)) + context_chunk_num = _uceildiv(context_length, fx.Int32(BKV)) + split_chunk_num = _uceildiv(context_chunk_num, split_kv) full_end = context_chunk_num * BKV split_cols = split_chunk_num * BKV tile_lo_col = pid_split_kv * split_cols - tile_hi_col = arith.minsi(_to_raw(tile_lo_col + split_cols), _to_raw(full_end)) + tile_hi_col = _imin(tile_lo_col + split_cols, full_end) ctx_m1 = context_length - 1 - for iv in scf.for_( - _to_raw(fx.Index(tile_lo_col)), - _to_raw(fx.Index(fx.Int32(tile_hi_col))), - _to_raw(fx.Index(fx.Int32(BKV))), - ): - col0 = fx.Int32(arith.index_cast(T.i32, iv)) + for col0 in range(tile_lo_col, tile_hi_col, fx.Int32(BKV)): wave_ni_base = wave * N_TILES_PER_WAVE # Helpers close over this chunk's SSA values; default args silence B023. @@ -242,9 +239,7 @@ def _col_for(abs_ni, _col0=col0): return _col0 + abs_ni * MFMA_N + lane_mod_N def _col_c_for(abs_ni, _ctx_m1=ctx_m1): - return fx.Int32( - arith.minsi(_to_raw(_col_for(abs_ni)), _to_raw(_ctx_m1)) - ) + return _imin(_col_for(abs_ni), _ctx_m1) def _ind_off_for(abs_ni): return pid_batch * max_block_len + udiv(_col_c_for(abs_ni), KVB) @@ -256,14 +251,9 @@ def _load_physical(ind_off): return fx.Int32(ind_t[ind_off]) def _prefetch_physical(physical_cur, ind_off_cur, ind_off_next): - same = _to_raw(ind_off_next) == _to_raw(ind_off_cur) - phys_raw = _to_raw(physical_cur) - if_op = scf.IfOp(same, results_=[T.i32], has_else=True) - with ir.InsertionPoint(if_op.then_block): - scf.YieldOp([phys_raw]) - with ir.InsertionPoint(if_op.else_block): - scf.YieldOp([_to_raw(_load_physical(ind_off_next))]) - return fx.Int32(if_op.results[0]) + if ind_off_next == ind_off_cur: + return physical_cur + return _load_physical(ind_off_next) def _b_col_for_tile(physical, tok_in_block, b0_carried): return [ @@ -275,18 +265,21 @@ def _process_tile_with_b0(abs_ni, physical, b0_carried): col = _col_for(abs_ni) tok_in_block = _tok_in_block_for(abs_ni) b_col = _b_col_for_tile(physical, tok_in_block, b0_carried) - kv_scale = _to_raw(fx.Float32(kv_scale_v[physical, tok_in_block])) - if scale_mul != 1.0: - kv_scale = arith.MulFOp( - kv_scale, scale_mul_c, fastmath=fm_fast - ).result + kv_scale = fx.Float32(kv_scale_v[physical, tok_in_block]) col_sum = mfma_head_reduce( a_pack, b_col, w_frag, kv_scale, m_tiles=M_TILES, k_steps=K_STEPS ) - is_writer = _to_raw((lane_div_N == 0) & (col <= q_limit)) - with ir.InsertionPoint(scf.IfOp(is_writer).then_block): - out_row_t[col] = fx.Float32(col_sum) - scf.YieldOp([]) + is_writer = (lane_div_N == 0) & (col <= q_limit) + + def _do_write(_t=out_row_t, _c=col, _v=col_sum): + _t[_c] = _v + + @flyc.jit + def _guarded_write(_pred=is_writer, _w=_do_write): + if _pred: + _w() + + _guarded_write() def _prefetch_k0(abs_ni, physical): return load_pack_kv( @@ -393,14 +386,11 @@ def compile_fp8_paged_mqa_logits( kv_block_size: int = 1, preshuffle: bool = False, variant: str = DEFAULT_VARIANT, - scale_mul: float = 1.0, ): """Return a cached, compiled FlyDSL paged launcher for the given config. ``num_heads``/``head_size``/``kv_block_size``/``preshuffle`` are compile-time - constants and ``variant`` is a ``paged_w`` tag. ``scale_mul`` folds a - constant into the co-packed dequant scale; it carries the gfx942 FN->FNUZ - 2x compensation. + constants and ``variant`` is a ``paged_w`` tag. """ launcher = _build_paged_kernel( num_heads=num_heads, @@ -409,7 +399,6 @@ def compile_fp8_paged_mqa_logits( waves_per_block=_variant_wpb(variant), kv_block_size=kv_block_size, preshuffle=preshuffle, - scale_mul=scale_mul, ) launcher.compile_hints = dict(DEFAULT_COMPILE_HINTS) return launcher @@ -437,7 +426,7 @@ def flydsl_fp8_paged_mqa_logits( Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. - q_fp8: [batch, next_n, heads, hidden_dim], float8 e4m3 (fn or fnuz) + q_fp8: [batch, next_n, heads, hidden_dim], float8 e4m3fn (gfx950 native) kv_cache: [num_blocks, KVBlockSize, 1, index_dim] uint8, co-packed per block as KVBlockSize*hidden_dim fp8 K bytes then KVBlockSize f32 scales; index_dim == hidden_dim + 4 (+ optional padding). @@ -509,13 +498,10 @@ def flydsl_fp8_paged_mqa_logits( _fnuz, _fn, ), f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" + assert get_gfx() == "gfx950", ( + f"flydsl_fp8_paged_mqa_logits targets gfx950 (32x32x64 MFMA); got {get_gfx()}" + ) - # gfx942's fp8 MFMA reads operands as e4m3 FNUZ (bias 8), so an FN (OCP, - # bias 7) Q byte decodes to half its value there; scale_mul undoes that on - # the co-packed scale. NOTE: FN -0 (0x80) is FNUZ NaN and is *not* patched - # -- the vector<8xi32> operand has no byte-patch path. The KV cache is - # always arch-native fp8, so it never needs either. - scale_mul = 2.0 if (get_gfx() == "gfx942" and q_fp8.dtype == _fn) else 1.0 variant = _resolve_variant(variant) launcher = compile_fp8_paged_mqa_logits( @@ -525,7 +511,6 @@ def flydsl_fp8_paged_mqa_logits( kv_block_size=int(KVBlockSize), preshuffle=bool(Preshuffle), variant=variant, - scale_mul=scale_mul, ) # Co-packed cache -> raw uint8 byte view [num_blocks * KVBlockSize * index_dim]. diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index c5dd36df9c3..ad32e6440c6 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -1,11 +1,14 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness sweep for the FlyDSL paged FP8 MQA-logits (decode) kernel. +"""Correctness sweep and benchmarks for the FlyDSL paged FP8 MQA-logits kernel. The reference is a torch port of vLLM's ``fp8_paged_mqa_logits_torch``, copied here so the test does not depend on vLLM. Gate: exact ``-inf``-mask match plus ``calc_diff < 1e-3``, tolerances NOT widened. + +Performance timing and optional Gluon A/B comparison live in this file rather +than separate benchmark scripts. """ import argparse @@ -19,23 +22,18 @@ import aiter from aiter import dtypes from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from aiter.test_common import benchmark, checkAllclose +from aiter.test_common import benchmark, checkAllclose, run_perftest torch.set_default_device("cuda") -# gfx942 is the primary target; gfx950 (native fp8) is also runnable. -SUPPORTED_GFX = ["gfx942", "gfx950"] -# "fnuz" == the arch-native e4m3 (fnuz on gfx942, fn on gfx950); "fn" forces the -# OCP e4m3fn operand (exercises the gfx942 FN->FNUZ patch). Mirrors the dense -# test's dtype aliasing. +SUPPORTED_GFX = ["gfx950"] _E4M3_NATIVE = get_fp8_e4m3_dtype() DTYPE_MAP = {"fnuz": _E4M3_NATIVE, "fn": torch.float8_e4m3fn} -try: - from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits -except ImportError: - flydsl_fp8_paged_mqa_logits = None +REF_CHUNK_K = 256 +REF_WAVE_PER_EU = 2 def calc_diff(x, y): @@ -100,17 +98,10 @@ def ref_fp8_paged_mqa_logits( fp8_dtype, block_size=1, ): - """Torch reference, vectorized port of vLLM ``fp8_paged_mqa_logits_torch``:: - - logits[b*next_n+n, p] = sum_h ReLU() * weights[.., h] - - masked by ``p <= context_len - next_n + n``. Position ``p`` resolves to - ``(block_tables[b, p//KVB], p%KVB)`` in the block-flat co-packed cache. - """ + """Torch reference, vectorized port of vLLM ``fp8_paged_mqa_logits_torch``.""" batch_size, next_n, _heads, dim = q.size() num_blocks = kv_cache_fp8.shape[0] index_dim = kv_cache_fp8.shape[-1] - # Reconstruct the block-flat block: [KVB keys (dim bytes each)][KVB f32 scales]. flat = kv_cache_fp8.reshape(num_blocks, block_size * index_dim) keys = ( flat[:, : block_size * dim] @@ -125,7 +116,7 @@ def ref_fp8_paged_mqa_logits( .view(torch.float32) .view(num_blocks, block_size, 1) ) - kvf = keys * scales # [num_blocks, block_size, dim] dequantized K + kvf = keys * scales qf = q.float() logits = torch.full( [batch_size * next_n, max_model_len], @@ -138,27 +129,25 @@ def ref_fp8_paged_mqa_logits( if context_len == 0: continue pos = torch.arange(context_len, device=q.device) - blk = block_tables[i, pos // block_size] # [ctx] physical block ids - tok = pos % block_size # [ctx] token-in-block - kx = kvf[blk, tok] # [ctx, dim] dequantized K - s = torch.einsum("nhd,pd->nhp", qf[i], kx) # [next_n, heads, ctx] + blk = block_tables[i, pos // block_size] + tok = pos % block_size + kx = kvf[blk, tok] + s = torch.einsum("nhd,pd->nhp", qf[i], kx) s = torch.relu(s) - wl = weights[i * next_n : (i + 1) * next_n, :] # [next_n, heads] - s = (s * wl[:, :, None]).sum(dim=1) # [next_n, ctx] + wl = weights[i * next_n : (i + 1) * next_n, :] + s = (s * wl[:, :, None]).sum(dim=1) q_lim = ( context_len - next_n + torch.arange(next_n, device=q.device) - ).unsqueeze( - 1 - ) # [next_n, 1] + ).unsqueeze(1) s = torch.where(pos[None, :] <= q_lim, s, float("-inf")) logits[i * next_n : (i + 1) * next_n, :context_len] = s return logits class Inputs(NamedTuple): - q: torch.Tensor # bf16, for the oracle - q_fp8: torch.Tensor # quantized, for the kernel - kv_cache_fp8: torch.Tensor # co-packed, unshuffled + q: torch.Tensor + q_fp8: torch.Tensor + kv_cache_fp8: torch.Tensor weights: torch.Tensor context_lens: torch.Tensor block_tables: torch.Tensor @@ -167,31 +156,41 @@ class Inputs(NamedTuple): def _build_inputs( - batch_size, next_n, heads, head_dim, avg_kv_length, q_dtype, block_size=1, seed=0 + batch_size, + next_n, + heads, + head_dim, + avg_kv_length, + q_dtype, + block_size=1, + seed=0, + var_ratio=0.0, + pool_blocks=0, ): torch.manual_seed(seed) random.seed(seed) fp8_dtype = get_fp8_e4m3_dtype() max_model_len = 2 * avg_kv_length - num_blocks = (max_model_len + block_size - 1) // block_size - - lo = max(1, int((1 - 0.5) * avg_kv_length)) - hi = int((1 + 0.5) * avg_kv_length) + 1 - context_lens = torch.randint(lo, hi, (batch_size,)).cuda().to(torch.int32) - # decode with MTP needs at least next_n tokens of context. + if var_ratio == 0.0: + context_lens = torch.full( + (batch_size,), avg_kv_length, device="cuda", dtype=torch.int32 + ) + else: + lo = max(1, int((1 - var_ratio) * avg_kv_length)) + hi = int((1 + var_ratio) * avg_kv_length) + 1 + context_lens = torch.randint(lo, hi, (batch_size,)).cuda().to(torch.int32) context_lens = torch.clamp(context_lens, min=next_n) + blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size + max_block_len = int(blocks_per_seq.max().item()) + needed_blocks = int(blocks_per_seq.sum().item()) + num_blocks = needed_blocks if pool_blocks <= 0 else max(pool_blocks, max_block_len) + q = torch.randn((batch_size, next_n, heads, head_dim), dtype=torch.bfloat16) kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) weights = torch.randn((batch_size * next_n, heads), dtype=torch.float32) - # Block table: one entry per KVBlockSize-token block, handed out of a - # shuffled pool in sequence order. The pool is sized from max_model_len, so - # sequences share blocks once the batch needs more than it holds. That is - # deliberate: it costs nothing for correctness and exercises block reuse. - blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size - max_block_len = int(blocks_per_seq.max().item()) pool = list(range(num_blocks)) random.shuffle(pool) pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) @@ -217,31 +216,35 @@ def _build_inputs( ) -@benchmark() -def test_fp8_paged_mqa_logits( +def _kernel_inputs(inp, batch_size, next_n, head_dim, preshuffle, block_size): + kv_cache_kernel = ( + preshuffle_kv_data(inp.kv_cache_fp8, head_dim) + if preshuffle + else inp.kv_cache_fp8 + ) + out = torch.full( + (batch_size * next_n, inp.max_model_len), + float("-inf"), + device="cuda", + dtype=torch.float32, + ) + return kv_cache_kernel, out + + +def _verify_paged_mqa_logits( + inp, batch_size, next_n, - heads, head_dim, - avg_kv_length, - q_dtype, - split_kv=0, - block_size=1, - preshuffle=False, - variant=None, - chunk_k=128, + split_kv, + block_size, + preshuffle, + variant, + chunk_k, ): - # split_kv == 0 -> auto; else an explicit override (1 disables splitting). - # Both must be correctness-identical. _split_kv = None if split_kv == 0 else split_kv - inp = _build_inputs( - batch_size, - next_n, - heads, - head_dim, - avg_kv_length, - DTYPE_MAP[q_dtype], - block_size=block_size, + kv_cache_kernel, out = _kernel_inputs( + inp, batch_size, next_n, head_dim, preshuffle, block_size ) with torch.inference_mode(): @@ -257,20 +260,6 @@ def test_fp8_paged_mqa_logits( ) ref_mask = ref == float("-inf") - out = torch.full( - (batch_size * next_n, inp.max_model_len), - float("-inf"), - device="cuda", - dtype=torch.float32, - ) - - # Same numbers as the oracle's cache, in the production data layout. - kv_cache_kernel = ( - preshuffle_kv_data(inp.kv_cache_fp8, head_dim) - if preshuffle - else inp.kv_cache_fp8 - ) - with torch.inference_mode(): got = flydsl_fp8_paged_mqa_logits( inp.q_fp8, @@ -303,8 +292,83 @@ def test_fp8_paged_mqa_logits( msg="flydsl paged fp8_mqa_logits", printLog=False, ) + return err, kv_cache_kernel, out + + +@benchmark() +def _bench_flydsl_paged_kernel( + q_fp8, + kv_cache_kernel, + weights, + out, + context_lens, + block_tables, + max_model_len, + *, + preshuffle, + block_size, + split_kv, + chunk_k, + variant, +): + _split_kv = None if split_kv == 0 else split_kv + + def fn(): + return flydsl_fp8_paged_mqa_logits( + q_fp8, + kv_cache_kernel, + weights, + out, + context_lens, + block_tables, + max_model_len, + Preshuffle=preshuffle, + KVBlockSize=block_size, + SplitKV=_split_kv, + ChunkK=chunk_k, + variant=variant, + ) - return { + with torch.inference_mode(): + _, us = run_perftest(fn) + return {"flydsl us": float(us)} + + +def test_fp8_paged_mqa_logits( + batch_size, + next_n, + heads, + head_dim, + avg_kv_length, + q_dtype, + split_kv=0, + block_size=1, + preshuffle=False, + variant=None, + chunk_k=128, + bench=False, +): + inp = _build_inputs( + batch_size, + next_n, + heads, + head_dim, + avg_kv_length, + DTYPE_MAP[q_dtype], + block_size=block_size, + ) + err, kv_cache_kernel, out = _verify_paged_mqa_logits( + inp, + batch_size, + next_n, + head_dim, + split_kv, + block_size, + preshuffle, + variant, + chunk_k, + ) + ret = { "gfx": get_gfx(), "kvb": block_size, "preshuffle": preshuffle, @@ -313,17 +377,58 @@ def test_fp8_paged_mqa_logits( "chunk_k": chunk_k, "flydsl err": err, } + if not bench: + return ret + + ret.update( + _bench_flydsl_paged_kernel( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + preshuffle=preshuffle, + block_size=block_size, + split_kv=split_kv, + chunk_k=chunk_k, + variant=variant, + ) + ) + + if block_size == 1 and not preshuffle: + from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits + + out_ref = torch.full( + (batch_size * next_n, inp.max_model_len), + float("-inf"), + device="cuda", + dtype=torch.float32, + ) + + def fn_ref(): + return deepgemm_fp8_paged_mqa_logits( + inp.q_fp8, + inp.kv_cache_fp8, + inp.weights, + out_ref, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + ChunkK=REF_CHUNK_K, + Preshuffle=False, + KVBlockSize=1, + WavePerEU=REF_WAVE_PER_EU, + ) + + with torch.inference_mode(): + _, ref_us = run_perftest(fn_ref) + ret["gluon us"] = float(ref_us) + + return ret -# The kernel specialises at compile time on (heads, head_dim, KVBlockSize, -# Preshuffle); batch, next_n, context length, SplitKV and the Q dtype are -# runtime values. Those axes do not interact -- nothing couples head_dim to -# SplitKV, or the dtype to KVBlockSize -- so crossing all seven produced 864 -# cases that compiled only 12 distinct kernels and re-ran each about 72 times. -# -# Instead: cross the compile-time axes, vary each runtime axis alone from a -# base, and add back the few combinations where the tile-range arithmetic -# genuinely couples. `--exhaustive` still runs the full product. _BASE = { "batch_size": 2, "next_n": 2, @@ -333,8 +438,6 @@ def test_fp8_paged_mqa_logits( "q_dtype": "fnuz", } _SHAPES = [(64, 64), (64, 128), (128, 64), (128, 128)] - - _AXES = ( "batch_size", "next_n", @@ -348,14 +451,11 @@ def test_fp8_paged_mqa_logits( def _c(**kw): - """A case: the base config with KVBlockSize=64, plus the overrides given.""" return {**_BASE, "block_size": 64, **kw} def default_cases(): cases = [] - # Compile-time cross: one case per distinct compiled kernel. Preshuffle - # needs KVBlockSize % 16 == 0. for heads, head_dim in _SHAPES: for kvb in (1, 64): cases.append(_c(heads=heads, head_dim=head_dim, block_size=kvb)) @@ -363,20 +463,13 @@ def default_cases(): cases.append( _c(heads=heads, head_dim=head_dim, block_size=kvb, preshuffle=True) ) - # Runtime axes, one at a time off the base. cases += [_c(batch_size=b, next_n=n) for b, n in ((1, 1), (1, 2), (4, 2), (8, 1))] - cases += [_c(avg_kv_length=kv) for kv in (128, 8192)] # 128 < one ChunkK tile - cases += [_c(split_kv=sk) for sk in (1, 4)] # 1 disables splitting - cases.append(_c(q_dtype="fn")) # gfx942 FN->FNUZ Q patch - # Interactions: SplitKV against the context length it partitions, and - # against the preshuffled layout, which the old matrix never crossed. + cases += [_c(avg_kv_length=kv) for kv in (128, 8192)] + cases += [_c(split_kv=sk) for sk in (1, 4)] cases.append(_c(avg_kv_length=128, split_kv=4)) cases.append(_c(avg_kv_length=8192, split_kv=1)) cases.append(_c(split_kv=4, preshuffle=True)) cases.append(_c(split_kv=1, block_size=16, preshuffle=True)) - # Wave split and KV tile width, neither covered before: a bug in the - # wave-to-column-tile assignment is correct by construction at one wave per - # CTA. ChunkK/32 must be divisible by the variant's wave count. cases += [_c(variant=v) for v in ("paged_w2", "paged_w4")] cases.append(_c(chunk_k=64, variant="paged_w2")) cases.append(_c(chunk_k=256, variant="paged_w4")) @@ -386,13 +479,12 @@ def default_cases(): def exhaustive_cases(): - """The original 7-way cartesian product (870 cases), for release validation.""" prod = itertools.product( [(1, 1), (1, 2), (2, 1), (2, 2), (4, 2), (8, 1)], [64, 128], [64, 128], [128, 1024, 8192], - ["fnuz", "fn"], + ["fnuz"], [0, 1, 4], [1, 64], ) @@ -400,7 +492,6 @@ def exhaustive_cases(): dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, sk, kvb))) for (bs, nn), nh, hd, kv, qd, sk, kvb in prod ] - # Preshuffle was covered by a hand-picked set rather than a full cross. return cases + [ {**dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, 0, kvb))), "preshuffle": True} for bs, nn, nh, hd, kv, qd, kvb in [ @@ -408,46 +499,196 @@ def exhaustive_cases(): (1, 2, 64, 128, 1024, "fnuz", 64), (2, 1, 128, 128, 8192, "fnuz", 64), (1, 1, 64, 64, 1024, "fnuz", 16), - (1, 1, 64, 128, 1024, "fn", 64), (4, 2, 64, 128, 8192, "fnuz", 64), ] ] +def gluon_ab_shapes(): + H, D = 64, 128 + shapes = [] + for B in (1, 4, 16, 64, 128): + for avg_kv in (16384, 32768, 65536): + shapes.append((B, 2, H, D, avg_kv)) + for B in (1, 16, 128): + shapes.append((B, 1, H, D, 32768)) + return shapes + + +def run_profile(args): + inp = _build_inputs( + args.batch, + args.next_n, + args.heads, + args.head_dim, + args.kv_len, + DTYPE_MAP[args.q_dtype], + block_size=args.kv_block_size, + seed=args.seed, + var_ratio=args.var_ratio, + pool_blocks=args.pool_blocks, + ) + err, kv_cache_kernel, out = _verify_paged_mqa_logits( + inp, + args.batch, + args.next_n, + args.head_dim, + args.split_kv, + args.kv_block_size, + args.preshuffle, + args.variant, + args.chunk_k, + ) + _split_kv = None if args.split_kv == 0 else args.split_kv + + def run(): + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + Preshuffle=args.preshuffle, + KVBlockSize=args.kv_block_size, + ChunkK=args.chunk_k, + SplitKV=_split_kv, + WavePerEU=args.wave_per_eu, + variant=args.variant, + ) + + run() + torch.cuda.synchronize() + total_ctx = int(inp.context_lens.sum().item()) + kv_bytes = args.next_n * total_ctx * (args.head_dim + 4) + print(f"# profile correctness err={err:.3e}") + print( + f"# gfx={get_gfx()} B={args.batch} nn={args.next_n} H={args.heads} " + f"D={args.head_dim} kv_len={args.kv_len} kvb={args.kv_block_size} " + f"preshuffle={args.preshuffle} chunk_k={args.chunk_k} iters={args.iters}" + ) + print(f"# KV requested: {kv_bytes / 1e6:.1f} MB") + torch.cuda.synchronize() + for _ in range(args.iters): + run() + torch.cuda.synchronize() + if args.time: + _, us = run_perftest(run, num_iters=args.num_iters) + flops = 2 * args.heads * args.head_dim * args.next_n * total_ctx + print( + f"time: {us:.2f} us | {flops / us / 1e6:.1f} TFLOP/s | " + f"{kv_bytes / us / 1e3:.1f} GB/s requested" + ) + + +def run_gluon_ab(args): + rows = [] + for B, nn, H, D, kv_len in gluon_ab_shapes(): + ret = test_fp8_paged_mqa_logits( + batch_size=B, + next_n=nn, + heads=H, + head_dim=D, + avg_kv_length=kv_len, + q_dtype="fnuz", + block_size=1, + preshuffle=False, + bench=True, + ) + fly_us = ret["flydsl us"] + ref_us = ret.get("gluon us", float("nan")) + rows.append( + { + **ret, + "B": B, + "nn": nn, + "avg_kv": kv_len, + "fly/ref": fly_us / ref_us if ref_us == ref_us and ref_us > 0 else float("nan"), + } + ) + print( + f"B={B} nn={nn} kv={kv_len} fly={fly_us:.1f}us gluon={ref_us:.1f}us " + f"fly/ref={rows[-1]['fly/ref']:.2f}x err={ret['flydsl err']:.2e}", + flush=True, + ) + df = pd.DataFrame(rows) + try: + summary = df.to_markdown(index=False) + except ImportError: + summary = df.to_string(index=False) + aiter.logger.info("gluon A/B summary:\n%s", summary) + + def main(): if get_gfx() not in SUPPORTED_GFX: aiter.logger.warning( "fp8_paged_mqa_logits unsupported on %s; skipping", get_gfx() ) return - if flydsl_fp8_paged_mqa_logits is None: - aiter.logger.warning("flydsl package not installed; skipping") - return parser = argparse.ArgumentParser( - description="FlyDSL paged fp8_mqa_logits correctness sweep" + description="FlyDSL paged fp8_mqa_logits correctness + benchmarks" ) parser.add_argument( "--exhaustive", action="store_true", - help="run the full 7-way cartesian product (870 cases) instead of the " - "curated matrix; for release validation", + help="run the full cartesian product instead of the curated matrix", ) + parser.add_argument("--no-preshuffle", action="store_true", help="skip preshuffle cases") parser.add_argument( - "--no-preshuffle", action="store_true", help="skip the preshuffle cases" + "--bench", + action="store_true", + help="also run run_perftest timing in the correctness sweep", + ) + parser.add_argument( + "--compare-gluon", + action="store_true", + help="run the production Gluon A/B decode shape sweep (KVBlockSize=1)", + ) + parser.add_argument( + "--profile", + action="store_true", + help="run a single-config timed loop (for rocprofv3 / isolation profiling)", ) + parser.add_argument("--batch", type=int, default=16) + parser.add_argument("--next-n", type=int, default=2) + parser.add_argument("--heads", type=int, default=64) + parser.add_argument("--head-dim", type=int, default=128) + parser.add_argument("--kv-len", type=int, default=32768) + parser.add_argument("--kv-block-size", type=int, default=1) + parser.add_argument("--preshuffle", action="store_true") + parser.add_argument("--q-dtype", type=str, default="fnuz", choices=["fnuz", "fn"]) + parser.add_argument("--split-kv", type=int, default=0) + parser.add_argument("--wave-per-eu", type=int, default=2) + parser.add_argument("--chunk-k", type=int, default=128) + parser.add_argument("--variant", type=str, default=None) + parser.add_argument("--var-ratio", type=float, default=0.0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--pool-blocks", type=int, default=0) + parser.add_argument("--iters", type=int, default=100) + parser.add_argument("--num-iters", type=int, default=101) + parser.add_argument("--time", action="store_true") args = parser.parse_args() + if args.profile: + if args.preshuffle and args.kv_block_size % 16 != 0: + raise SystemExit("--preshuffle requires --kv-block-size divisible by 16") + run_profile(args) + return + if args.compare_gluon: + run_gluon_ab(args) + return + cases = exhaustive_cases() if args.exhaustive else default_cases() if args.no_preshuffle: cases = [c for c in cases if not c.get("preshuffle")] - df = [test_fp8_paged_mqa_logits(**c) for c in cases] + df = [test_fp8_paged_mqa_logits(bench=args.bench, **c) for c in cases] df = pd.DataFrame(df) try: summary = df.to_markdown(index=False) except ImportError: - # `tabulate` (optional) not installed -> plain-text table. summary = df.to_string(index=False) aiter.logger.info("fp8_paged_mqa_logits summary:\n%s", summary) diff --git a/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py b/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py deleted file mode 100644 index 311516a2ee9..00000000000 --- a/op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py +++ /dev/null @@ -1,335 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Standalone microbenchmark for the FlyDSL paged FP8 MQA-logits (decode) kernel. - -Runs ONE config of ``flydsl_fp8_paged_mqa_logits`` in isolation (no reference / -Gluon comparison) so the kernel can be profiled directly. Example rocprofv3 run:: - - PYTHONPATH=. rocprofv3 --kernel-trace --output-format csv \\ - --output-directory .rocprofv3 -- \\ - python op_tests/op_benchmarks/flydsl/bench_flydsl_paged_mqa_logits.py \\ - --heads 64 --head-dim 128 --kv-len 32768 --batch 16 --next-n 2 \\ - --kv-block-size 64 --preshuffle --iters 200 --no-check - -Covers both the block-flat gather (``--kv-block-size N``, default 1) and the -production preshuffle layout (``--preshuffle``; requires KVBlockSize % 16 == 0). -The input build, co-pack byte layout, preshuffle, and torch oracle are imported -from the flydsl correctness test, so the profiled kernel sees exactly the tested -layout (single source of truth). - -By default the run is correctness-gated once (torch oracle, exact ``-inf`` mask + -calc_diff < 1e-3) before the timed loop; pass ``--no-check`` to skip it (handy -when profiling). The timed loop launches the kernel ``--iters`` times back to -back -- that is what rocprofv3 traces. ``--time`` additionally reports device -self-time via ``run_perftest``. -""" - -import argparse -import random - -import torch - -from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits -from aiter.ops.flydsl.kernels.mqa_logits.fp8_paged_mqa_logits import DEFAULT_VARIANT -from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from aiter.test_common import run_perftest - -# Reuse the correctness test's layout/oracle helpers so the profiled kernel sees -# exactly the tested co-pack + preshuffle byte layout. -from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( - calc_diff, - kv_cache_cast_to_fp8, - preshuffle_kv_data, - ref_fp8_paged_mqa_logits, -) - -torch.set_default_device("cuda") - -DTYPE_MAP = {"fnuz": get_fp8_e4m3_dtype(), "fn": torch.float8_e4m3fn} - - -def build_inputs( - batch, - next_n, - heads, - head_dim, - kv_len, - block_size, - q_dtype, - var_ratio, - seed, - pool_blocks=0, -): - """Paged decode inputs (block-flat co-packed cache), block_size >= 1. - - ``var_ratio`` controls the per-sequence context-length spread around - ``kv_len``: 0.0 == exact (every sequence == kv_len); v>0 draws lengths in - [(1-v), (1+v)]*kv_len. A fixed ``seed`` makes the whole build deterministic - (stable kernel work across profiling runs). - - The pool holds one distinct physical block per block the batch actually - needs, so the KV footprint the kernel touches is ``sum(ctx)*index_dim`` -- - what a real decode reads. Sizing the pool from ``max_model_len`` instead - (as an earlier version did) makes the block table wrap: at batch=16, - kv_len=32768, KVBlockSize=64 that is a 1024-block pool serving 8192 block - references, so all 16 sequences share the same blocks and the working set - collapses from 69 MB to 8.65 MB -- small enough to sit in cache and make - the kernel look ~13% faster than it is. ``pool_blocks`` overrides the size - to study cache residency deliberately (values below the demand re-introduce - wrapping; the padded tail of each row stays block 0 either way). - """ - torch.manual_seed(seed) - random.seed(seed) - fp8_dtype = get_fp8_e4m3_dtype() - - max_model_len = 2 * kv_len - - if var_ratio == 0.0: - context_lens = torch.full((batch,), kv_len, device="cuda", dtype=torch.int32) - else: - lo = max(1, int((1 - var_ratio) * kv_len)) - hi = int((1 + var_ratio) * kv_len) + 1 - context_lens = torch.randint(lo, hi, (batch,)).cuda().to(torch.int32) - # decode with MTP needs at least next_n tokens of context. - context_lens = torch.clamp(context_lens, min=next_n) - - # block table: ceil(ctx / block_size) physical blocks per sequence. - blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size - max_block_len = int(blocks_per_seq.max().item()) - needed_blocks = int(blocks_per_seq.sum().item()) - num_blocks = needed_blocks if pool_blocks <= 0 else max(pool_blocks, max_block_len) - - q = torch.randn((batch, next_n, heads, head_dim), dtype=torch.bfloat16) - kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) - weights = torch.randn((batch * next_n, heads), dtype=torch.float32) - - # Hand blocks out of a shuffled pool in sequence order (so a sequence's - # blocks are scattered across the pool, as they are after real paged - # allocation), padding each row's unused tail with block 0. - pool = list(range(num_blocks)) - random.shuffle(pool) - pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) - starts = torch.cumsum(blocks_per_seq, 0) - blocks_per_seq - col = torch.arange(max_block_len, device="cuda", dtype=torch.int64) - in_row = col[None, :] < blocks_per_seq[:, None] - draw = (starts[:, None] + col[None, :]) % num_blocks - block_tables = torch.where( - in_row, pool_t[draw], torch.zeros((), device="cuda", dtype=torch.int32) - ).to(torch.int32) - - q_fp8 = q.to(q_dtype) - kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, fp8_dtype) - return ( - q, - q_fp8, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - ) - - -def main(): - parser = argparse.ArgumentParser( - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - description=__doc__, - ) - # Single-config by design: profile exactly one kernel launch shape. - parser.add_argument("--batch", type=int, default=16) - parser.add_argument("--next-n", type=int, default=2) - parser.add_argument("--heads", type=int, default=64) - parser.add_argument("--head-dim", type=int, default=128) - parser.add_argument("--kv-len", type=int, default=32768) - parser.add_argument( - "--kv-block-size", type=int, default=1, help="KVBlockSize (block-flat page)" - ) - parser.add_argument( - "--preshuffle", - action="store_true", - help="use the shuffle_weight(16x16) KV layout (needs kv-block-size %% 16 == 0)", - ) - parser.add_argument("--q-dtype", type=str, default="fnuz", choices=["fnuz", "fn"]) - parser.add_argument( - "--split-kv", - type=int, - default=0, - help="0 == auto (host formula); else override", - ) - parser.add_argument("--wave-per-eu", type=int, default=2) - parser.add_argument("--chunk-k", type=int, default=128, help="KV tile width") - parser.add_argument( - "--variant", - type=str, - default=None, - help=f"paged_w tag; None uses the kernel default ({DEFAULT_VARIANT})", - ) - parser.add_argument("--var-ratio", type=float, default=0.0) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument( - "--pool-blocks", - type=int, - default=0, - help="physical blocks in the cache pool; 0 == one per block the batch " - "needs (no sharing). Smaller values make sequences share blocks, " - "shrinking the footprint into cache", - ) - parser.add_argument( - "--iters", type=int, default=100, help="kernel launches in the timed loop" - ) - parser.add_argument( - "--num-iters", type=int, default=101, help="run_perftest iters for --time" - ) - parser.add_argument( - "--time", - action="store_true", - help="also report device self-time (run_perftest)", - ) - parser.add_argument( - "--no-check", action="store_true", help="skip the one-shot correctness gate" - ) - args = parser.parse_args() - - if get_gfx() not in ("gfx942", "gfx950"): - print(f"unsupported gfx {get_gfx()}; skipping") - return - if args.preshuffle and args.kv_block_size % 16 != 0: - raise SystemExit( - f"--preshuffle requires --kv-block-size divisible by 16; " - f"got {args.kv_block_size}" - ) - - ( - q, - q_fp8, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - ) = build_inputs( - args.batch, - args.next_n, - args.heads, - args.head_dim, - args.kv_len, - args.kv_block_size, - DTYPE_MAP[args.q_dtype], - args.var_ratio, - args.seed, - args.pool_blocks, - ) - - # Oracle reads the unshuffled cache; the kernel gets the preshuffled copy. - kv_kernel = ( - preshuffle_kv_data(kv_cache_fp8, args.head_dim) - if args.preshuffle - else kv_cache_fp8 - ) - - out = torch.full( - (args.batch * args.next_n, max_model_len), - float("-inf"), - device="cuda", - dtype=torch.float32, - ) - split_kv = None if args.split_kv == 0 else args.split_kv - - def run(): - flydsl_fp8_paged_mqa_logits( - q_fp8, - kv_kernel, - weights, - out, - context_lens, - block_tables, - max_model_len, - Preshuffle=args.preshuffle, - KVBlockSize=args.kv_block_size, - ChunkK=args.chunk_k, - SplitKV=split_kv, - WavePerEU=args.wave_per_eu, - variant=args.variant, - ) - - # Warm / JIT-compile once (excluded from the timed loop and from rocprofv3 - # runs that filter on the steady-state kernel name). - run() - torch.cuda.synchronize() - - if not args.no_check: - ref = ref_fp8_paged_mqa_logits( - q, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - block_size=args.kv_block_size, - ) - neg_inf = float("-inf") - ref_mask = ref == neg_inf - diff = float( - calc_diff(out.masked_fill(out == neg_inf, 0), ref.masked_fill(ref_mask, 0)) - ) - mask_ok = bool(torch.equal(out == neg_inf, ref_mask)) - ok = (diff < 1e-3) and mask_ok - print( - f"# correctness: calc_diff={diff:.3e} mask_ok={mask_ok} " - f"-> {'PASS' if ok else 'FAIL'}" - ) - if not ok: - raise SystemExit("correctness gate failed") - - # Work estimate: QK dot is 2*H*D MACs per (query row, valid KV position). - total_ctx = int(context_lens.sum().item()) - flops = 2 * args.heads * args.head_dim * args.next_n * total_ctx - # Bytes the kernel *requests*: every query row walks its whole context, so - # the next_n rows of a sequence each re-read the same blocks. - kv_bytes = args.next_n * total_ctx * (args.head_dim + 4) - # Bytes the kernel *touches*: distinct blocks the table points at. Equal to - # kv_bytes/next_n unless the pool is small enough that sequences share. - n_blk_row = (context_lens.to(torch.int64) + args.kv_block_size - 1) // ( - args.kv_block_size - ) - in_row = ( - torch.arange(block_tables.shape[1], device="cuda")[None, :] < n_blk_row[:, None] - ) - distinct_blocks = int(torch.unique(block_tables[in_row]).numel()) - kv_unique = distinct_blocks * args.kv_block_size * (args.head_dim + 4) - - print( - f"# gfx={get_gfx()} B={args.batch} nn={args.next_n} H={args.heads} " - f"D={args.head_dim} kv_len={args.kv_len} kvb={args.kv_block_size} " - f"preshuffle={args.preshuffle} chunk_k={args.chunk_k} " - f"split_kv={args.split_kv or 'auto'} variant={args.variant or DEFAULT_VARIANT} " - f"q_dtype={args.q_dtype} iters={args.iters}" - ) - print( - f"# KV footprint: {kv_unique / 1e6:.1f} MB touched " - f"({distinct_blocks} blocks), {kv_bytes / 1e6:.1f} MB requested" - ) - - # Timed loop: back-to-back kernel launches -- this is what rocprofv3 traces. - torch.cuda.synchronize() - for _ in range(args.iters): - run() - torch.cuda.synchronize() - - if args.time: - _, us = run_perftest(run, num_iters=args.num_iters) - us = float(us) - print( - f"time: {us:.2f} us | {flops / us / 1e6:.1f} TFLOP/s | " - f"{kv_bytes / us / 1e3:.1f} GB/s requested | " - f"{kv_unique / us / 1e3:.1f} GB/s touched" - ) - - -if __name__ == "__main__": - main() diff --git a/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py b/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py deleted file mode 100644 index 2e381b20daa..00000000000 --- a/op_tests/op_benchmarks/flydsl/bench_flydsl_vs_gluon_paged_mqa_logits.py +++ /dev/null @@ -1,427 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Honest A/B: FlyDSL paged MQA-logits vs the production Gluon/Triton kernel. - -Both kernels are timed by the SAME code path: identical inputs (built once and -fed to both), identical ``run_perftest`` methodology (device self-time, IQR -cleaning, same warmup/iters), identical output dtype/layout, and both are -correctness-gated against the SAME torch reference before any timing is reported. - -The reference kernel is the production ``deepgemm_fp8_paged_mqa_logits`` host -(aiter Triton/Gluon), run with vLLM's production defaults (ChunkK=256, -WavePerEU=2, KVBlockSize=1, Preshuffle=False) -- NOT detuned. Both the reference -host and the FlyDSL host auto-compute SplitKV to fill the device on small decode -grids, so the comparison is apples-to-apples. - -This A/B is fixed at ``KVBlockSize=1``: the Gluon reference faults on a -non-padded block-flat cache at ``KVBlockSize>1``, so any KVBlockSize>1 (incl. -Preshuffle) comparison would not be apples-to-apples. FlyDSL KVBlockSize>1 and -Preshuffle correctness is covered against a torch oracle in -``op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py`` instead. - -Reuses the paged benchmark's fp8 cache packer to guarantee the same input build. -""" - -import argparse -import random -import statistics - -import torch - -from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits -from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits -from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from aiter.test_common import run_perftest - -# same fp8 co-pack builder the production paged benchmark uses -from op_tests.op_benchmarks.triton.bench_deepgemm_attention import ( - kv_cache_cast_to_fp8, -) - -torch.set_default_device("cuda") - -# Production reference config (matches vLLM rocm_fp8_paged_mqa_logits for decode). -REF_CHUNK_K = 256 -REF_WAVE_PER_EU = 2 - - -def calc_diff(x, y): - x, y = x.double(), y.double() - denominator = (x * x + y * y).sum() - return (1 - 2 * (x * y).sum() / denominator).item() - - -def ref_fp8_paged_mqa_logits( - q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype -): - """Vectorized torch reference (dequantizes the co-packed fp8 cache).""" - batch_size, next_n, _heads, dim = q.size() - kvv, scale = kv_cache_fp8[..., :dim], kv_cache_fp8[..., dim:] - scale = scale.contiguous().view(torch.float) - qf = q.float() - kvf = (kvv.view(fp8_dtype).float() * scale).view(kv_cache_fp8.shape[0], dim) - logits = torch.full( - [batch_size * next_n, max_model_len], - float("-inf"), - device=q.device, - dtype=torch.float32, - ) - for i in range(batch_size): - context_len = int(context_lens[i].item()) - if context_len == 0: - continue - pages = block_tables[i, :context_len] - kx = kvf[pages] - s = torch.einsum("nhd,pd->nhp", qf[i], kx) - s = torch.relu(s) - wl = weights[i * next_n : (i + 1) * next_n, :] - s = (s * wl[:, :, None]).sum(dim=1) - p = torch.arange(context_len, device=q.device) - q_lim = ( - context_len - next_n + torch.arange(next_n, device=q.device) - ).unsqueeze(1) - s = torch.where(p[None, :] <= q_lim, s, float("-inf")) - logits[i * next_n : (i + 1) * next_n, :context_len] = s - return logits - - -def parity_mask(context_lens, batch_size, next_n, max_model_len): - """Causal-valid mask (bench methodology): positions <= ctx - next_n + off.""" - positions = ( - torch.arange(max_model_len, device="cuda") - .unsqueeze(0) - .expand(batch_size * next_n, -1) - ) - row_indices = torch.arange(batch_size * next_n, device="cuda") // next_n - next_n_offset = torch.arange(batch_size * next_n, device="cuda") % next_n - return positions <= (context_lens[row_indices] - next_n + next_n_offset).unsqueeze( - 1 - ) - - -def build_inputs( - batch_size, - next_n, - heads, - index_dim, - kv_length, - seed=0, - var_ratio=0.0, - pool_blocks=0, -): - """Paged input builder (blocksize==1). - - ``var_ratio`` controls the per-sequence context-length spread around - ``kv_length``: 0.0 gives an exact length (every sequence == kv_length), and - v>0 draws lengths uniformly from [(1-v)*kv_length, (1+v)*kv_length]. - - The pool holds one distinct block per block the batch needs, so the KV - footprint both kernels read is ``sum(ctx)*(index_dim+4)`` -- what a real - decode touches. Sizing it from ``max_model_len`` instead (as an earlier - version did) makes the block table wrap: at blocksize=1 the pool is - 2*kv_length blocks while the batch needs batch_size*kv_length, so from - batch_size=4 up the sequences share blocks and the working set collapses - into cache. At B=16, kv_length=32768 that is 8.65 MB standing in for 69 MB, - which flatters both kernels and, because the two spend their time - differently, does not flatter them equally. ``pool_blocks`` overrides the - size to study cache residency deliberately. - """ - torch.manual_seed(seed) - random.seed(seed) - fp8_dtype = get_fp8_e4m3_dtype() - - max_model_len = 2 * kv_length - blocksize = 1 - - if var_ratio == 0.0: - context_lens = torch.full( - (batch_size,), kv_length, device="cuda", dtype=torch.int32 - ) - else: - context_lens = ( - torch.randint( - int((1 - var_ratio) * kv_length), - int((1 + var_ratio) * kv_length) + 1, - (batch_size,), - ) - .cuda() - .to(torch.int32) - ) - # decode with MTP needs at least next_n tokens of context. - context_lens = torch.clamp(context_lens, min=next_n) - - blocks_per_seq = (context_lens.to(torch.int64) + blocksize - 1) // blocksize - max_block_len = int(blocks_per_seq.max().item()) - needed_blocks = int(blocks_per_seq.sum().item()) - num_blocks = needed_blocks if pool_blocks <= 0 else max(pool_blocks, max_block_len) - - # The kernels ride the per-token byte offset on a 32-bit buffer voffset, so - # the whole pool must be addressable in i32 bytes. - pool_bytes = num_blocks * blocksize * (index_dim + 4) - if pool_bytes >= 2**31: - raise SystemExit( - f"pool of {num_blocks} blocks = {pool_bytes / 1e9:.2f} GB exceeds the " - f"i32 gather-offset limit (2.15 GB). Shrink the shape or cap the pool " - f"with --pool-blocks (which re-introduces block sharing)." - ) - - q = torch.randn( - (batch_size, next_n, heads, index_dim), device="cuda", dtype=torch.bfloat16 - ) - kv_cache = torch.randn( - (num_blocks, blocksize, 1, index_dim), device="cuda", dtype=torch.bfloat16 - ) - weights = torch.randn( - (batch_size * next_n, heads), device="cuda", dtype=torch.float32 - ) - - # Hand blocks out of a shuffled pool in sequence order (scattered across the - # pool, as they are after real paged allocation); pad each row's tail with 0. - pool = list(range(num_blocks)) - random.shuffle(pool) - pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) - starts = torch.cumsum(blocks_per_seq, 0) - blocks_per_seq - col = torch.arange(max_block_len, device="cuda", dtype=torch.int64) - in_row = col[None, :] < blocks_per_seq[:, None] - draw = (starts[:, None] + col[None, :]) % num_blocks - block_tables = torch.where( - in_row, pool_t[draw], torch.zeros((), device="cuda", dtype=torch.int32) - ).to(torch.int32) - - q_fp8 = q.to(fp8_dtype) - kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, padding=False, fp8_dtype=fp8_dtype) - return ( - q, - q_fp8, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - min(needed_blocks, num_blocks) * blocksize * (index_dim + 4), - ) - - -def time_kernel(fn, repeats, num_iters): - us = [] - for _ in range(repeats): - _, u = run_perftest(fn, num_iters=num_iters) - us.append(float(u)) - return statistics.median(us), min(us), max(us) - - -def run_shape( - batch_size, - next_n, - heads, - index_dim, - kv_length, - repeats, - num_iters, - var_ratio=0.0, - pool_blocks=0, -): - ( - q, - q_fp8, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - touched_bytes, - ) = build_inputs( - batch_size, - next_n, - heads, - index_dim, - kv_length, - var_ratio=var_ratio, - pool_blocks=pool_blocks, - ) - - ref = ref_fp8_paged_mqa_logits( - q, kv_cache_fp8, weights, context_lens, block_tables, max_model_len, fp8_dtype - ) - valid = parity_mask(context_lens, batch_size, next_n, max_model_len) - neg_inf = float("-inf") - ref_neg_mask = ref == neg_inf - - def correctness(out): - diff = calc_diff(out.masked_fill(~valid, 0), ref.masked_fill(~valid, 0)) - mask_ok = bool(torch.equal(out == neg_inf, ref_neg_mask)) - return diff, mask_ok - - # ---- reference (production Gluon/Triton), production defaults ---- - out_ref = torch.full( - (batch_size * next_n, max_model_len), - neg_inf, - device="cuda", - dtype=torch.float32, - ) - - def fn_ref(): - deepgemm_fp8_paged_mqa_logits( - q_fp8, - kv_cache_fp8, - weights, - out_ref, - context_lens, - block_tables, - max_model_len, - ChunkK=REF_CHUNK_K, - Preshuffle=False, - KVBlockSize=1, - WavePerEU=REF_WAVE_PER_EU, - ) - - # ---- FlyDSL (KVBlockSize==1, auto SplitKV) ---- - out_fly = torch.full( - (batch_size * next_n, max_model_len), - neg_inf, - device="cuda", - dtype=torch.float32, - ) - - def fn_fly(): - flydsl_fp8_paged_mqa_logits( - q_fp8, - kv_cache_fp8, - weights, - out_fly, - context_lens, - block_tables, - max_model_len, - ) - - # Warm / JIT-compile once (excluded from the reported device self-time anyway). - fn_ref() - fn_fly() - torch.cuda.synchronize() - - ref_diff, ref_mask_ok = correctness(out_ref) - fly_diff, fly_mask_ok = correctness(out_fly) - ref_pass = (ref_diff < 1e-3) and ref_mask_ok - fly_pass = (fly_diff < 1e-3) and fly_mask_ok - - ref_med = ref_min = ref_max = fly_med = fly_min = fly_max = float("nan") - if ref_pass and fly_pass: - ref_med, ref_min, ref_max = time_kernel(fn_ref, repeats, num_iters) - fly_med, fly_min, fly_max = time_kernel(fn_fly, repeats, num_iters) - - return { - "B": batch_size, - "nn": next_n, - "grid": batch_size * next_n, - "H": heads, - "D": index_dim, - "avg_kv": kv_length, - "kv_mb": touched_bytes / 1e6, - "ref_ms": ref_med / 1e3, - "ref_min_ms": ref_min / 1e3, - "ref_max_ms": ref_max / 1e3, - "fly_ms": fly_med / 1e3, - "fly_min_ms": fly_min / 1e3, - "fly_max_ms": fly_max / 1e3, - "fly/ref": ( - (fly_med / ref_med) - if ref_med == ref_med and ref_med > 0 # noqa: PLR0124 (NaN check) - else float("nan") - ), - "ref_diff": ref_diff, - "fly_diff": fly_diff, - "ref_mask": ref_mask_ok, - "fly_mask": fly_mask_ok, - "ref_pass": ref_pass, - "fly_pass": fly_pass, - } - - -# DSA lightning-indexer decode dims (heads=64, index_dim=128) swept over long -# contexts and a concurrency (batch) range; next_n=2 is speculative decode, -# next_n=1 non-speculative. All shapes reported, wins and losses. -def default_shapes(): - H, D = 64, 128 - shapes = [] - for B in (1, 4, 16, 64, 128): - for avg_kv in (16384, 32768, 65536): - shapes.append((B, 2, H, D, avg_kv)) - for B in (1, 16, 128): - shapes.append((B, 1, H, D, 32768)) - return shapes - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--repeats", type=int, default=5) - parser.add_argument("--num-iters", type=int, default=101) - parser.add_argument( - "--var-ratio", - type=float, - default=0.0, - help="Per-sequence context-length spread around kv_len. 0.0 = exact " - "(every sequence == kv_len); e.g. 0.5 draws lengths in [0.5, 1.5]*kv_len.", - ) - parser.add_argument( - "--pool-blocks", - type=int, - default=0, - help="physical blocks in the cache pool; 0 == one per block the batch " - "needs (no sharing). Smaller values make sequences share blocks, " - "shrinking the KV footprint into cache", - ) - args = parser.parse_args() - - if get_gfx() not in ("gfx942", "gfx950"): - print(f"unsupported gfx {get_gfx()}; skipping") - return - - len_mode = "exact" if args.var_ratio == 0.0 else f"uniform +/-{args.var_ratio:g}" - print( - f"# arch={get_gfx()} fp8={get_fp8_e4m3_dtype()} " - f"repeats={args.repeats} num_iters={args.num_iters} " - f"ref_chunk_k={REF_CHUNK_K} ref_wave_per_eu={REF_WAVE_PER_EU} " - f"ctx_len={len_mode} max_model_len=2*kv_len blocksize=1 " - f"pool={'one block per block needed' if args.pool_blocks <= 0 else args.pool_blocks}" - ) - print( - "# kv_mb = distinct KV bytes the batch touches " - "(both kernels read it once per query row)" - ) - header = ( - "B nn grid H D kv_len kv_mb | ref_ms[min-max] fly_ms[min-max] fly/ref | " - "ref_diff fly_diff ref_mask fly_mask ref_pass fly_pass" - ) - print(header) - rows = [] - for B, nn, H, D, kv_len in default_shapes(): - r = run_shape( - B, - nn, - H, - D, - kv_len, - args.repeats, - args.num_iters, - var_ratio=args.var_ratio, - pool_blocks=args.pool_blocks, - ) - rows.append(r) - print( - f"{r['B']} {r['nn']} {r['grid']} {r['H']} {r['D']} {r['avg_kv']} " - f"{r['kv_mb']:.0f} | " - f"{r['ref_ms']:.4f}[{r['ref_min_ms']:.4f}-{r['ref_max_ms']:.4f}] " - f"{r['fly_ms']:.4f}[{r['fly_min_ms']:.4f}-{r['fly_max_ms']:.4f}] " - f"{r['fly/ref']:.2f}x | " - f"{r['ref_diff']:.2e} {r['fly_diff']:.2e} " - f"{r['ref_mask']} {r['fly_mask']} {r['ref_pass']} {r['fly_pass']}", - flush=True, - ) - - -if __name__ == "__main__": - main() From f2bc4d1cf44ccd2ad3ad61a78a52fe030fc69ba7 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 25 Aug 2026 12:07:44 +0000 Subject: [PATCH 29/52] style(flydsl): fix Black/Ruff on paged MQA logits files Apply Black formatting and resolve Ruff PLR0124/I001 so CI pre-checks pass. Co-authored-by: Cursor --- .../kernels/mqa_logits/_mqa_logits_common.py | 2 +- .../kernels/mqa_logits/fp8_paged_mqa_logits.py | 6 +++--- .../test_flydsl_fp8_paged_mqa_logits.py | 15 ++++++++++++--- 3 files changed, 16 insertions(+), 7 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py index d6577443e18..d8eff26e26d 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py @@ -3,7 +3,7 @@ """FlyDSL building blocks for the paged FP8 MQA-logits kernel (gfx950 / CDNA4). - logits[row, n] = sum_h ReLU( * kv_scale[n]) * weights[row, h] +logits[row, n] = sum_h ReLU( * kv_scale[n]) * weights[row, h] """ from functools import lru_cache diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 242a164546d..95bf6326c70 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -498,9 +498,9 @@ def flydsl_fp8_paged_mqa_logits( _fnuz, _fn, ), f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" - assert get_gfx() == "gfx950", ( - f"flydsl_fp8_paged_mqa_logits targets gfx950 (32x32x64 MFMA); got {get_gfx()}" - ) + assert ( + get_gfx() == "gfx950" + ), f"flydsl_fp8_paged_mqa_logits targets gfx950 (32x32x64 MFMA); got {get_gfx()}" variant = _resolve_variant(variant) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index ad32e6440c6..a36e812fa4a 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -13,6 +13,7 @@ import argparse import itertools +import math import random from typing import NamedTuple @@ -398,7 +399,9 @@ def test_fp8_paged_mqa_logits( ) if block_size == 1 and not preshuffle: - from aiter.ops.triton.attention.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits + from aiter.ops.triton.attention.pa_mqa_logits import ( + deepgemm_fp8_paged_mqa_logits, + ) out_ref = torch.full( (batch_size * next_n, inp.max_model_len), @@ -604,7 +607,11 @@ def run_gluon_ab(args): "B": B, "nn": nn, "avg_kv": kv_len, - "fly/ref": fly_us / ref_us if ref_us == ref_us and ref_us > 0 else float("nan"), + "fly/ref": ( + fly_us / ref_us + if not math.isnan(ref_us) and ref_us > 0 + else float("nan") + ), } ) print( @@ -635,7 +642,9 @@ def main(): action="store_true", help="run the full cartesian product instead of the curated matrix", ) - parser.add_argument("--no-preshuffle", action="store_true", help="skip preshuffle cases") + parser.add_argument( + "--no-preshuffle", action="store_true", help="skip preshuffle cases" + ) parser.add_argument( "--bench", action="store_true", From f33a573e74f7fef06d4f10686c4c30dcc57c330e Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Sun, 6 Sep 2026 11:38:55 +0000 Subject: [PATCH 30/52] feat(flydsl): add gfx950 H64 16x16x128 paged FP8 MQA logits kernel Land the production indexer mapping (resident Q, depth-4 NT B-ring, SplitKV) without the v2/v3 experiments or local profiling artifacts. Co-authored-by: Cursor --- aiter/ops/flydsl/__init__.py | 4 + .../_fp8_paged_mqa_logits_gfx950.py | 154 ++++++++ .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 352 ++++++++++++++++++ ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 164 ++++++++ 4 files changed, 674 insertions(+) create mode 100644 aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py create mode 100644 aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py create mode 100644 op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 92a202e1bb8..123db663e1b 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -57,6 +57,9 @@ from .kernels.mqa_logits.fp8_paged_mqa_logits import ( flydsl_fp8_paged_mqa_logits, ) + from .kernels.mqa_logits.fp8_paged_mqa_logits_gfx950 import ( + flydsl_fp8_paged_mqa_logits_gfx950, + ) from .kernels.mqa_logits.pa_mqa_logits_fp4 import ( flydsl_pa_mqa_logits_fp4, ) @@ -80,6 +83,7 @@ "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", "flydsl_fp8_paged_mqa_logits", + "flydsl_fp8_paged_mqa_logits_gfx950", "flydsl_hgemm", "flydsl_mla_reduce_v1", "flydsl_moe_stage1", diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py new file mode 100644 index 00000000000..f4a5292c5ec --- /dev/null +++ b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""16x16x128 FP8 primitives for the gfx950 paged-MQA indexer fast path.""" + +import flydsl.compiler as flyc +import flydsl.expr as fx +from flydsl.expr import arith, range_constexpr, rocdl +from flydsl.expr.typing import T + +from .. import buffer_ops + +Vec = fx.Vector + +NUM_HEADS = 64 +HEAD_DIM = 128 +KV_BLOCK_SIZE = 64 +NEXT_N_MAX = 2 +MFMA_M = 16 +MFMA_N = 16 +MFMA_K = 128 +M_TILES = NUM_HEADS // MFMA_M +DREG = 4 +B_RING = KV_BLOCK_SIZE // MFMA_N +# Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). +PAGE_VMEM_LOADS = B_RING * 3 +# Predicated logit stores: 2 query rows × 4 tiles. Hot-loop pages are full, so +# these 8 VM ops land after the next-page loads and can stay in flight. +PAGE_VMEM_STORES = NEXT_N_MAX * B_RING + +_NEUTRAL_E8M0 = 0x7F7F7F7F + + +def wait_vmcnt(n): + """Explicit s_waitcnt vmcnt(n). Leaves n VM ops outstanding.""" + rocdl.sched_barrier(0) + rocdl.s_waitcnt(vmcnt=int(n)) + rocdl.sched_barrier(0) + + +def uceildiv(a, b): + a, b = fx.Int32(a), fx.Int32(b) + return fx.Int32((fx.Uint32(a) + fx.Uint32(b) - 1) // fx.Uint32(b)) + + +def imin(a, b): + a, b = fx.Int32(a), fx.Int32(b) + return (a <= b).select(a, b) + + +def guarded_store(pred, store_fn): + """Predicated side effect (FlyDSL ``@flyc.jit`` branch on a runtime mask).""" + + @flyc.jit + def _guarded(_pred=pred, _store=store_fn): + if _pred: + _store() + + _guarded() + + +def _concat_i32x4(lo, hi): + lo, hi = Vec(lo), Vec(hi) + return Vec.from_elements( + [Vec(lo)[i].ir_value() for i in range_constexpr(4)] + + [Vec(hi)[i].ir_value() for i in range_constexpr(4)], + fx.Int32, + ) + + +def load_q_pack(q_i32, byte_base, lane_div_16): + """Load one row's lane-owned contiguous K32 segment as a 256-bit operand.""" + off = (byte_base + lane_div_16 * 32) // fx.Int32(4) + lo = buffer_ops.buffer_load(q_i32.rsrc, off, vec_width=4, dtype=T.i32) + hi = buffer_ops.buffer_load(q_i32.rsrc, off + 4, vec_width=4, dtype=T.i32) + return _concat_i32x4(lo, hi) + + +def load_preshuffled_k_pack( + kv_i32, + physical, + tile_in_page, + lane_mod_16, + lane_div_16, + *, + index_dim, +): + """Load a shuffle_weight(16,16) K column for 16x16x128 scaled MFMA. + + shuffle_weight stores ``[K/32, K-half, token16, byte16]``. ``lane_div_16`` + selects the lane's K32 segment and the two 16-byte halves form i32x8. + """ + block_i32 = fx.Int32(KV_BLOCK_SIZE * index_dim // 4) + lane_k_bytes = MFMA_K // (64 // MFMA_N) + base = ( + physical * block_i32 + + fx.Int32(tile_in_page * HEAD_DIM * MFMA_N // 4) + + lane_div_16 * fx.Int32(lane_k_bytes * MFMA_N // 4) + + lane_mod_16 * 4 + ) + lo = buffer_ops.buffer_load( + kv_i32.rsrc, base, vec_width=4, dtype=T.i32, cache_modifier=2 + ) + hi = buffer_ops.buffer_load( + kv_i32.rsrc, base + fx.Int32(MFMA_N * 16 // 4), + vec_width=4, dtype=T.i32, cache_modifier=2, + ) + return _concat_i32x4(lo, hi) + + +def load_kv_scale(kv_i32, physical, token_in_page, *, index_dim): + block_i32 = fx.Int32(KV_BLOCK_SIZE * index_dim // 4) + off = physical * block_i32 + fx.Int32(KV_BLOCK_SIZE * HEAD_DIM // 4) + token_in_page + return fx.Float32( + buffer_ops.buffer_load( + kv_i32.rsrc, off, vec_width=1, dtype=T.f32, cache_modifier=2 + ) + ) + + +def mfma_scores(a_tiles, b_pack): + """Return four 4-f32 score fragments, one per 16-head M tile.""" + result_type = Vec.make_type(DREG, fx.Float32) + neutral = arith.constant(_NEUTRAL_E8M0, type=T.i32) + scores = [] + for mi in range_constexpr(M_TILES): + acc = Vec.filled(DREG, 0.0, fx.Float32) + acc = rocdl.mfma_scale_f32_16x16x128_f8f6f4( + result_type, + [a_tiles[mi], b_pack, acc, 0, 0, 0, neutral, 0, neutral], + ) + scores.append(acc) + return scores + + +def reduce_scores(scores, weights, kv_scale): + """ReLU, weighted H reduction, positive KV scale, then wave reduction.""" + zero = fx.Float32(0.0) + total = zero + for mi in range_constexpr(M_TILES): + frag = Vec(scores[mi]) + for ii in range_constexpr(DREG): + total = total + fx.Float32(frag[ii]).maximumf(zero) * weights[mi][ii] + total = total * kv_scale + total = total + total.shuffle_xor(16, 64) + total = total + total.shuffle_xor(32, 64) + return total + + +def schedule_mfma_valu_pairs(): + """Pair each 32-cycle MFMA with the prior tile's 12-op VALU fragment.""" + for _ in range_constexpr(M_TILES): + rocdl.sched_group_barrier(0x008, 1, 0) + rocdl.sched_group_barrier(0x002, DREG * 3, 0) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py new file mode 100644 index 00000000000..9eed19a98c0 --- /dev/null +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -0,0 +1,352 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""gfx950 H64/D128/KVB64 preshuffled paged FP8 MQA-logits kernel.""" + +from functools import lru_cache + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch +from flydsl.expr import arith, range_constexpr, rocdl +from flydsl.expr.typing import T + +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype + +from .. import buffer_ops +from ..tensor_shim import GTensor, _run_compiled, _to_raw +from ._fp8_paged_mqa_logits_gfx950 import ( + B_RING, + DREG, + HEAD_DIM, + MFMA_M, + MFMA_N, + M_TILES, + NEXT_N_MAX, + NUM_HEADS, + KV_BLOCK_SIZE, + guarded_store, + imin, + load_kv_scale, + load_preshuffled_k_pack, + load_q_pack, + mfma_scores, + reduce_scores, + schedule_mfma_valu_pairs, + uceildiv, + wait_vmcnt, + PAGE_VMEM_LOADS, + PAGE_VMEM_STORES, +) +from ._mqa_logits_common import ( + DEFAULT_COMPILE_HINTS, + device_cu_count, + udiv, + umod, +) + +_GFX950 = "gfx950" + + +def _unwrap(value): + return value.ir_value() if hasattr(value, "ir_value") else value + + +def _build_kernel(*, index_dim: int): + kernel_name = ( + f"fp8_paged_mqa_logits_gfx950_H{NUM_HEADS}_D{HEAD_DIM}_" + f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{NEXT_N_MAX}_w1_nt_ps_flydsl" + ) + + @flyc.kernel(name=kernel_name, known_block_size=[64, 1, 1]) + def kernel( + Q: fx.Tensor, + KV_cache: fx.Tensor, + weights: fx.Tensor, + out_logits: fx.Tensor, + context_lens: fx.Tensor, + kv_indices: fx.Tensor, + next_n: fx.Int32, + batch_size: fx.Int32, + split_kv: fx.Int32, + stride_q_batch: fx.Int32, + stride_q_next_n: fx.Int32, + stride_q_heads: fx.Int32, + max_block_len: fx.Int32, + stride_out: fx.Int32, + ): + bid = fx.block_idx.x + pid_batch = umod(bid, batch_size) + pid_split = udiv(bid, batch_size) + lane = umod(fx.thread_idx.x, 64) + lane_mod_16 = umod(lane, MFMA_N) + lane_div_16 = udiv(lane, MFMA_N) + + q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) + kv_i32 = GTensor(KV_cache, dtype=T.i32, shape=(-1,)) + weight_t = GTensor(weights, dtype=T.f32, shape=(-1,)) + context_t = GTensor(context_lens, dtype=T.i32, shape=(-1,)) + table_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) + out_t = GTensor(out_logits, dtype=T.f32, shape=(-1,), cache_modifier=2) + + context_len = fx.Int32(context_t[pid_batch]) + page_count = uceildiv(context_len, fx.Int32(KV_BLOCK_SIZE)) + pages_per_split = uceildiv(page_count, split_kv) + page_lo = pid_split * pages_per_split + page_hi = imin(page_lo + pages_per_split, page_count) + col_lo = page_lo * KV_BLOCK_SIZE + col_hi = page_hi * KV_BLOCK_SIZE + + a_rows = [ + [None for _ in range_constexpr(M_TILES)] + for _ in range_constexpr(NEXT_N_MAX) + ] + w_rows = [ + [[None for _ in range_constexpr(DREG)] for _ in range_constexpr(M_TILES)] + for _ in range_constexpr(NEXT_N_MAX) + ] + for row in range_constexpr(NEXT_N_MAX): + q_row = imin(fx.Int32(row), next_n - 1) + q_base = pid_batch * stride_q_batch + q_row * stride_q_next_n + out_row = pid_batch * next_n + q_row + for mi in range_constexpr(M_TILES): + h = mi * MFMA_M + lane_mod_16 + byte_base = q_base + h * stride_q_heads + a_rows[row][mi] = load_q_pack(q_i32, byte_base, lane_div_16) + weight_vec = fx.Vector( + buffer_ops.buffer_load( + weight_t.rsrc, + out_row * NUM_HEADS + mi * MFMA_M + lane_div_16 * DREG, + vec_width=DREG, + dtype=T.f32, + ) + ) + for ii in range_constexpr(DREG): + w_rows[row][mi][ii] = fx.Float32(weight_vec[ii]) + + def _issue_page(page_col): + table_idx = pid_batch * max_block_len + udiv(page_col, fx.Int32(KV_BLOCK_SIZE)) + physical = fx.Int32( + buffer_ops.buffer_load( + table_t.rsrc, table_idx, vec_width=1, dtype=T.i32, is_scalar=True + ) + ) + b_slots, scale_slots = [], [] + for slot in range_constexpr(B_RING): + token = slot * MFMA_N + lane_mod_16 + b_slots.append( + load_preshuffled_k_pack( + kv_i32, + physical, + slot, + lane_mod_16, + lane_div_16, + index_dim=index_dim, + ) + ) + scale_slots.append( + load_kv_scale(kv_i32, physical, token, index_dim=index_dim) + ) + return b_slots, scale_slots + + def _pack_state(b_slots, scale_slots): + return [_unwrap(v) for v in b_slots + scale_slots] + + def _unpack_state(state): + b_slots = [fx.Vector(state[i]) for i in range_constexpr(B_RING)] + scale_slots = [ + fx.Float32(state[B_RING + i]) for i in range_constexpr(B_RING) + ] + return b_slots, scale_slots + + def _store_row(row, col, scores, scale): + q_row = fx.Int32(row) + q_limit = context_len - next_n + q_row + out_row = pid_batch * next_n + q_row + logit = reduce_scores(scores, w_rows[row], scale) + writer = ( + (lane_div_16 == 0) + & (q_row < next_n) + & (col < context_len) + & (col <= q_limit) + ) + def _write(_row=out_row, _col=col, _value=logit): + out_t[_row * stride_out + _col] = _value + + guarded_store(writer, _write) + + def _compute_page(page_col, b_slots, scale_slots): + rocdl.sched_barrier(0) + prev_scores = None + prev_col = None + prev_scale = None + for slot in range_constexpr(B_RING): + col = page_col + slot * MFMA_N + lane_mod_16 + scores = [ + mfma_scores(a_rows[row], b_slots[slot]) + for row in range_constexpr(NEXT_N_MAX) + ] + if slot > 0: + schedule_mfma_valu_pairs() + for row in range_constexpr(NEXT_N_MAX): + _store_row(row, prev_col, prev_scores[row], prev_scale) + prev_scores = scores + prev_col = col + prev_scale = scale_slots[slot] + schedule_mfma_valu_pairs() + for row in range_constexpr(NEXT_N_MAX): + _store_row(row, prev_col, prev_scores[row], prev_scale) + rocdl.sched_barrier(0) + + if page_lo < page_hi: + b_init, scale_init = _issue_page(col_lo) + if page_lo + 1 < page_hi: + stop = col_hi - KV_BLOCK_SIZE + init_state = _pack_state(b_init, scale_init) + for page_col, state in range( + col_lo, stop, fx.Int32(KV_BLOCK_SIZE), init=init_state + ): + current_b, current_scale = _unpack_state(state) + next_b, next_scale = _issue_page( + fx.Int32(page_col) + fx.Int32(KV_BLOCK_SIZE) + ) + # 12 next-page loads in flight: wait until those 12 remain so + # the current page (or leftover stores) has retired. + wait_vmcnt(PAGE_VMEM_LOADS) + _compute_page(fx.Int32(page_col), current_b, current_scale) + wait_vmcnt(PAGE_VMEM_STORES) + results = yield _pack_state(next_b, next_scale) + final_b, final_scale = _unpack_state(results) + _compute_page(stop, final_b, final_scale) + else: + wait_vmcnt(0) + _compute_page(col_lo, b_init, scale_init) + + @flyc.jit + def launch( + Q, + KV_cache, + weights, + out_logits, + context_lens, + kv_indices, + grid_blocks, + next_n, + batch_size, + split_kv, + stride_q_batch, + stride_q_next_n, + stride_q_heads, + max_block_len, + stride_out, + stream: fx.Stream = fx.Stream(None), # noqa: B008 + ): + kernel._func.__name__ = kernel_name + kernel( + Q, + KV_cache, + weights, + out_logits, + context_lens, + kv_indices, + next_n, + batch_size, + split_kv, + stride_q_batch, + stride_q_next_n, + stride_q_heads, + max_block_len, + stride_out, + ).launch( + grid=(arith.index_cast(T.index, _to_raw(grid_blocks)), 1, 1), + block=(64, 1, 1), + stream=stream, + ) + + launch.compile_hints = dict(DEFAULT_COMPILE_HINTS) + launch.kernel_name = kernel_name + return launch + + +@lru_cache(maxsize=8) +def _compile(*, index_dim: int): + return _build_kernel(index_dim=index_dim) + + +def flydsl_fp8_paged_mqa_logits_gfx950( + q_fp8, + kv_cache, + weights, + out_logits, + context_lens, + kv_indices, + max_model_len, + *, + KVBlockSize=KV_BLOCK_SIZE, + SplitKV=None, + TotalCuCount=None, + stream=None, +): + """Run the production-only H64/D128/KVB64/preshuffle gfx950 mapping.""" + if get_gfx() != _GFX950: + raise RuntimeError(f"gfx950 kernel requested on {get_gfx()}") + batch_size, next_n, heads, head_dim = q_fp8.shape + if (heads, head_dim, int(KVBlockSize)) != (NUM_HEADS, HEAD_DIM, KV_BLOCK_SIZE): + raise ValueError("requires H=64, D=128, KVBlockSize=64") + if next_n not in (1, 2): + raise ValueError("requires next_n in {1, 2}") + if q_fp8.dtype != get_fp8_e4m3_dtype(): + raise ValueError(f"q_fp8 must be native FP8 E4M3, got {q_fp8.dtype}") + if kv_cache.dtype != torch.uint8: + raise ValueError("kv_cache must contain preshuffled uint8 FP8 data") + num_blocks, block_size, one, index_dim = kv_cache.shape + if block_size != KV_BLOCK_SIZE or one != 1 or index_dim != HEAD_DIM + 4: + raise ValueError(f"unexpected KV cache shape {tuple(kv_cache.shape)}") + + context_lens = context_lens.reshape(batch_size) + weights = weights.reshape(batch_size * next_n, NUM_HEADS) + max_block_len = kv_indices.shape[-1] + kv_indices = kv_indices.reshape(batch_size, max_block_len) + total_cu = ( + device_cu_count(q_fp8.device.index) + if TotalCuCount is None + else int(TotalCuCount) + ) + max_pages = max(1, (int(max_model_len) + KV_BLOCK_SIZE - 1) // KV_BLOCK_SIZE) + split_kv = ( + max( + 1, + min( + max_pages, + (total_cu * 4 * 2 + batch_size - 1) // batch_size, + ), + ) + if SplitKV is None or int(SplitKV) <= 0 + else max(1, min(max_pages, int(SplitKV))) + ) + grid_blocks = batch_size * split_kv + launcher = _compile(index_dim=int(index_dim)) + launcher.compile_hints = {**DEFAULT_COMPILE_HINTS, "waves_per_eu": 2} + stream = stream or torch.cuda.current_stream(q_fp8.device) + with torch.cuda.device(q_fp8.device.index): + _run_compiled( + launcher, + q_fp8, + kv_cache.reshape(-1), + weights, + out_logits, + context_lens, + kv_indices, + int(grid_blocks), + int(next_n), + int(batch_size), + int(split_kv), + int(q_fp8.stride(0)), + int(q_fp8.stride(1)), + int(q_fp8.stride(2)), + int(max_block_len), + int(out_logits.stride(0)), + stream, + ) + return out_logits diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py new file mode 100644 index 00000000000..061d8dc61b9 --- /dev/null +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Focused B16/Nq2/H64/D128/KVB64 test for the gfx950 indexer mapping.""" + +import argparse + +import torch + +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl import ( + flydsl_fp8_paged_mqa_logits, + flydsl_fp8_paged_mqa_logits_gfx950, +) +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from aiter.test_common import run_perftest +from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( + _build_inputs, + _kernel_inputs, + calc_diff, + ref_fp8_paged_mqa_logits, +) + +torch.set_default_device("cuda") + +BATCH = 16 +NEXT_N = 2 +HEADS = 64 +HEAD_DIM = 128 +KV_LEN = 32768 +KV_BLOCK_SIZE = 64 + + +def _inputs(): + inp = _build_inputs( + BATCH, + NEXT_N, + HEADS, + HEAD_DIM, + KV_LEN, + get_fp8_e4m3_dtype(), + block_size=KV_BLOCK_SIZE, + ) + kv_cache, out = _kernel_inputs( + inp, BATCH, NEXT_N, HEAD_DIM, True, KV_BLOCK_SIZE + ) + return inp, kv_cache, out + + +def _launch(inp, kv_cache, out): + return flydsl_fp8_paged_mqa_logits_gfx950( + inp.q_fp8, + kv_cache, + inp.weights, + out, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + KVBlockSize=KV_BLOCK_SIZE, + ) + + +def test_gfx950_indexer_mapping(): + assert get_gfx() == "gfx950" + inp, kv_cache, out = _inputs() + with torch.inference_mode(): + ref = ref_fp8_paged_mqa_logits( + inp.q, + inp.kv_cache_fp8, + inp.weights, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + inp.fp8_dtype, + block_size=KV_BLOCK_SIZE, + ) + got = _launch(inp, kv_cache, out) + + ref_mask = ref == float("-inf") + got_mask = got == float("-inf") + assert torch.equal(got_mask, ref_mask), "causal/padding -inf mask mismatch" + diff = calc_diff(got.masked_fill(got_mask, 0), ref.masked_fill(ref_mask, 0)) + assert diff < 1e-3, f"calc_diff={diff}" + print(f"correctness: pass calc_diff={float(diff):.3e}") + + +def _benchmark(): + from aiter.ops.triton.attention.pa_mqa_logits import ( + deepgemm_fp8_paged_mqa_logits, + enable_gluon_pa_mqa_logits, + triton_version, + ) + + inp, kv_cache, out_new = _inputs() + out_old = torch.full_like(out_new, float("-inf")) + out_triton = torch.full_like(out_new, float("-inf")) + + def new_kernel(): + return _launch(inp, kv_cache, out_new) + + def old_kernel(): + return flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache, + inp.weights, + out_old, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + ) + + def triton_kernel(): + return deepgemm_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache, + inp.weights, + out_triton, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + ChunkK=256, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + WavePerEU=2, + ) + + backend = "gluon" if enable_gluon_pa_mqa_logits else "triton-jit" + print(f"triton {triton_version} backend={backend}") + with torch.inference_mode(): + _, new_us = run_perftest(new_kernel, num_iters=50, num_warmup=8) + _, old_us = run_perftest(old_kernel, num_iters=50, num_warmup=8) + _, triton_us = run_perftest(triton_kernel, num_iters=50, num_warmup=8) + print(f"time: {new_us:.3f} us") + print(f"flydsl baseline: {old_us:.3f} us") + print(f"{backend}: {triton_us:.3f} us") + print(f"vs flydsl: {old_us / new_us:.3f}x") + print(f"vs {backend}: {triton_us / new_us:.3f}x") + + +def _profile(): + inp, kv_cache, out = _inputs() + with torch.inference_mode(): + _launch(inp, kv_cache, out) + torch.cuda.synchronize() + print("profile launch: pass") + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--profile", action="store_true") + parser.add_argument("--time", action="store_true") + args = parser.parse_args() + if args.profile: + _profile() + elif args.time: + _benchmark() + else: + test_gfx950_indexer_mapping() + + +if __name__ == "__main__": + main() From a811a1dc3231a2a025fca474ad38bf5fe03c1f40 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Mon, 7 Sep 2026 12:24:37 +0000 Subject: [PATCH 31/52] feat(flydsl): specialize gfx950 paged MQA logits for H=32 and H=64 Compile a separate 16x16x128 kernel per head count so the indexer matches gluon's ChunkQ=heads set without changing the FlyDSL page pipeline. Co-authored-by: Cursor --- .../_fp8_paged_mqa_logits_gfx950.py | 17 ++-- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 47 +++++----- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 90 +++++++++++++------ 3 files changed, 97 insertions(+), 57 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py index f4a5292c5ec..5a4207b2b38 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py @@ -12,14 +12,13 @@ Vec = fx.Vector -NUM_HEADS = 64 +SUPPORTED_HEADS = (32, 64) HEAD_DIM = 128 KV_BLOCK_SIZE = 64 NEXT_N_MAX = 2 MFMA_M = 16 MFMA_N = 16 MFMA_K = 128 -M_TILES = NUM_HEADS // MFMA_M DREG = 4 B_RING = KV_BLOCK_SIZE // MFMA_N # Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). @@ -118,12 +117,12 @@ def load_kv_scale(kv_i32, physical, token_in_page, *, index_dim): ) -def mfma_scores(a_tiles, b_pack): - """Return four 4-f32 score fragments, one per 16-head M tile.""" +def mfma_scores(a_tiles, b_pack, *, m_tiles): + """Return one 4-f32 score fragment per 16-head M tile.""" result_type = Vec.make_type(DREG, fx.Float32) neutral = arith.constant(_NEUTRAL_E8M0, type=T.i32) scores = [] - for mi in range_constexpr(M_TILES): + for mi in range_constexpr(m_tiles): acc = Vec.filled(DREG, 0.0, fx.Float32) acc = rocdl.mfma_scale_f32_16x16x128_f8f6f4( result_type, @@ -133,11 +132,11 @@ def mfma_scores(a_tiles, b_pack): return scores -def reduce_scores(scores, weights, kv_scale): +def reduce_scores(scores, weights, kv_scale, *, m_tiles): """ReLU, weighted H reduction, positive KV scale, then wave reduction.""" zero = fx.Float32(0.0) total = zero - for mi in range_constexpr(M_TILES): + for mi in range_constexpr(m_tiles): frag = Vec(scores[mi]) for ii in range_constexpr(DREG): total = total + fx.Float32(frag[ii]).maximumf(zero) * weights[mi][ii] @@ -147,8 +146,8 @@ def reduce_scores(scores, weights, kv_scale): return total -def schedule_mfma_valu_pairs(): +def schedule_mfma_valu_pairs(*, m_tiles): """Pair each 32-cycle MFMA with the prior tile's 12-op VALU fragment.""" - for _ in range_constexpr(M_TILES): + for _ in range_constexpr(m_tiles): rocdl.sched_group_barrier(0x008, 1, 0) rocdl.sched_group_barrier(0x002, DREG * 3, 0) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 9eed19a98c0..51469404fd4 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx950 H64/D128/KVB64 preshuffled paged FP8 MQA-logits kernel.""" +"""gfx950 H32/H64 D128/KVB64 preshuffled paged FP8 MQA-logits kernel.""" from functools import lru_cache @@ -22,10 +22,11 @@ HEAD_DIM, MFMA_M, MFMA_N, - M_TILES, NEXT_N_MAX, - NUM_HEADS, KV_BLOCK_SIZE, + PAGE_VMEM_LOADS, + PAGE_VMEM_STORES, + SUPPORTED_HEADS, guarded_store, imin, load_kv_scale, @@ -36,8 +37,6 @@ schedule_mfma_valu_pairs, uceildiv, wait_vmcnt, - PAGE_VMEM_LOADS, - PAGE_VMEM_STORES, ) from ._mqa_logits_common import ( DEFAULT_COMPILE_HINTS, @@ -53,9 +52,10 @@ def _unwrap(value): return value.ir_value() if hasattr(value, "ir_value") else value -def _build_kernel(*, index_dim: int): +def _build_kernel(*, index_dim: int, num_heads: int): + m_tiles = num_heads // MFMA_M kernel_name = ( - f"fp8_paged_mqa_logits_gfx950_H{NUM_HEADS}_D{HEAD_DIM}_" + f"fp8_paged_mqa_logits_gfx950_H{num_heads}_D{HEAD_DIM}_" f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{NEXT_N_MAX}_w1_nt_ps_flydsl" ) @@ -99,25 +99,25 @@ def kernel( col_hi = page_hi * KV_BLOCK_SIZE a_rows = [ - [None for _ in range_constexpr(M_TILES)] + [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(NEXT_N_MAX) ] w_rows = [ - [[None for _ in range_constexpr(DREG)] for _ in range_constexpr(M_TILES)] + [[None for _ in range_constexpr(DREG)] for _ in range_constexpr(m_tiles)] for _ in range_constexpr(NEXT_N_MAX) ] for row in range_constexpr(NEXT_N_MAX): q_row = imin(fx.Int32(row), next_n - 1) q_base = pid_batch * stride_q_batch + q_row * stride_q_next_n out_row = pid_batch * next_n + q_row - for mi in range_constexpr(M_TILES): + for mi in range_constexpr(m_tiles): h = mi * MFMA_M + lane_mod_16 byte_base = q_base + h * stride_q_heads a_rows[row][mi] = load_q_pack(q_i32, byte_base, lane_div_16) weight_vec = fx.Vector( buffer_ops.buffer_load( weight_t.rsrc, - out_row * NUM_HEADS + mi * MFMA_M + lane_div_16 * DREG, + out_row * num_heads + mi * MFMA_M + lane_div_16 * DREG, vec_width=DREG, dtype=T.f32, ) @@ -164,7 +164,7 @@ def _store_row(row, col, scores, scale): q_row = fx.Int32(row) q_limit = context_len - next_n + q_row out_row = pid_batch * next_n + q_row - logit = reduce_scores(scores, w_rows[row], scale) + logit = reduce_scores(scores, w_rows[row], scale, m_tiles=m_tiles) writer = ( (lane_div_16 == 0) & (q_row < next_n) @@ -184,17 +184,17 @@ def _compute_page(page_col, b_slots, scale_slots): for slot in range_constexpr(B_RING): col = page_col + slot * MFMA_N + lane_mod_16 scores = [ - mfma_scores(a_rows[row], b_slots[slot]) + mfma_scores(a_rows[row], b_slots[slot], m_tiles=m_tiles) for row in range_constexpr(NEXT_N_MAX) ] if slot > 0: - schedule_mfma_valu_pairs() + schedule_mfma_valu_pairs(m_tiles=m_tiles) for row in range_constexpr(NEXT_N_MAX): _store_row(row, prev_col, prev_scores[row], prev_scale) prev_scores = scores prev_col = col prev_scale = scale_slots[slot] - schedule_mfma_valu_pairs() + schedule_mfma_valu_pairs(m_tiles=m_tiles) for row in range_constexpr(NEXT_N_MAX): _store_row(row, prev_col, prev_scores[row], prev_scale) rocdl.sched_barrier(0) @@ -270,8 +270,8 @@ def launch( @lru_cache(maxsize=8) -def _compile(*, index_dim: int): - return _build_kernel(index_dim=index_dim) +def _compile(*, index_dim: int, num_heads: int): + return _build_kernel(index_dim=index_dim, num_heads=num_heads) def flydsl_fp8_paged_mqa_logits_gfx950( @@ -288,12 +288,15 @@ def flydsl_fp8_paged_mqa_logits_gfx950( TotalCuCount=None, stream=None, ): - """Run the production-only H64/D128/KVB64/preshuffle gfx950 mapping.""" + """Run the H32/H64 D128/KVB64/preshuffle gfx950 mapping.""" if get_gfx() != _GFX950: raise RuntimeError(f"gfx950 kernel requested on {get_gfx()}") batch_size, next_n, heads, head_dim = q_fp8.shape - if (heads, head_dim, int(KVBlockSize)) != (NUM_HEADS, HEAD_DIM, KV_BLOCK_SIZE): - raise ValueError("requires H=64, D=128, KVBlockSize=64") + if heads not in SUPPORTED_HEADS or (head_dim, int(KVBlockSize)) != ( + HEAD_DIM, + KV_BLOCK_SIZE, + ): + raise ValueError("requires H in {32, 64}, D=128, KVBlockSize=64") if next_n not in (1, 2): raise ValueError("requires next_n in {1, 2}") if q_fp8.dtype != get_fp8_e4m3_dtype(): @@ -305,7 +308,7 @@ def flydsl_fp8_paged_mqa_logits_gfx950( raise ValueError(f"unexpected KV cache shape {tuple(kv_cache.shape)}") context_lens = context_lens.reshape(batch_size) - weights = weights.reshape(batch_size * next_n, NUM_HEADS) + weights = weights.reshape(batch_size * next_n, heads) max_block_len = kv_indices.shape[-1] kv_indices = kv_indices.reshape(batch_size, max_block_len) total_cu = ( @@ -326,7 +329,7 @@ def flydsl_fp8_paged_mqa_logits_gfx950( else max(1, min(max_pages, int(SplitKV))) ) grid_blocks = batch_size * split_kv - launcher = _compile(index_dim=int(index_dim)) + launcher = _compile(index_dim=int(index_dim), num_heads=int(heads)) launcher.compile_hints = {**DEFAULT_COMPILE_HINTS, "waves_per_eu": 2} stream = stream or torch.cuda.current_stream(q_fp8.device) with torch.cuda.device(q_fp8.device.index): diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index 061d8dc61b9..e86612b5a27 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Focused B16/Nq2/H64/D128/KVB64 test for the gfx950 indexer mapping.""" +"""Focused B16/Nq2/H{32,64}/D128/KVB64 test for the gfx950 indexer mapping.""" import argparse @@ -25,24 +25,25 @@ BATCH = 16 NEXT_N = 2 -HEADS = 64 +HEADS = (32, 64) +DEFAULT_HEADS = 64 HEAD_DIM = 128 KV_LEN = 32768 KV_BLOCK_SIZE = 64 -def _inputs(): +def _inputs(heads, next_n=NEXT_N, kv_len=KV_LEN, batch=BATCH): inp = _build_inputs( - BATCH, - NEXT_N, - HEADS, + batch, + next_n, + heads, HEAD_DIM, - KV_LEN, + kv_len, get_fp8_e4m3_dtype(), block_size=KV_BLOCK_SIZE, ) kv_cache, out = _kernel_inputs( - inp, BATCH, NEXT_N, HEAD_DIM, True, KV_BLOCK_SIZE + inp, batch, next_n, HEAD_DIM, True, KV_BLOCK_SIZE ) return inp, kv_cache, out @@ -60,9 +61,7 @@ def _launch(inp, kv_cache, out): ) -def test_gfx950_indexer_mapping(): - assert get_gfx() == "gfx950" - inp, kv_cache, out = _inputs() +def _check(inp, kv_cache, out, tag): with torch.inference_mode(): ref = ref_fp8_paged_mqa_logits( inp.q, @@ -78,20 +77,35 @@ def test_gfx950_indexer_mapping(): ref_mask = ref == float("-inf") got_mask = got == float("-inf") - assert torch.equal(got_mask, ref_mask), "causal/padding -inf mask mismatch" + assert torch.equal(got_mask, ref_mask), f"{tag}: causal/padding -inf mask mismatch" diff = calc_diff(got.masked_fill(got_mask, 0), ref.masked_fill(ref_mask, 0)) - assert diff < 1e-3, f"calc_diff={diff}" - print(f"correctness: pass calc_diff={float(diff):.3e}") + assert diff < 1e-3, f"{tag} calc_diff={diff}" + print(f"correctness {tag}: pass calc_diff={float(diff):.3e}") + +def test_gfx950_indexer_mapping(heads=DEFAULT_HEADS): + assert get_gfx() == "gfx950" + inp, kv_cache, out = _inputs(heads) + _check(inp, kv_cache, out, f"H={heads}") -def _benchmark(): + +def test_gfx950_nq1_and_short_pages(heads=32): + assert get_gfx() == "gfx950" + inp, kv_cache, out = _inputs(heads, next_n=1) + _check(inp, kv_cache, out, f"H={heads} Nq=1") + for pages in (2, 3, 4): + inp, kv_cache, out = _inputs(heads, kv_len=pages * KV_BLOCK_SIZE, batch=2) + _check(inp, kv_cache, out, f"H={heads} pages={pages}") + + +def _benchmark(heads): from aiter.ops.triton.attention.pa_mqa_logits import ( deepgemm_fp8_paged_mqa_logits, enable_gluon_pa_mqa_logits, triton_version, ) - inp, kv_cache, out_new = _inputs() + inp, kv_cache, out_new = _inputs(heads) out_old = torch.full_like(out_new, float("-inf")) out_triton = torch.full_like(out_new, float("-inf")) @@ -127,37 +141,61 @@ def triton_kernel(): ) backend = "gluon" if enable_gluon_pa_mqa_logits else "triton-jit" - print(f"triton {triton_version} backend={backend}") + print(f"H={heads} triton {triton_version} backend={backend}") with torch.inference_mode(): _, new_us = run_perftest(new_kernel, num_iters=50, num_warmup=8) _, old_us = run_perftest(old_kernel, num_iters=50, num_warmup=8) _, triton_us = run_perftest(triton_kernel, num_iters=50, num_warmup=8) - print(f"time: {new_us:.3f} us") + # WaveScope benchmarkPattern matches the first `time: {us} us` line (H=64). + if heads == DEFAULT_HEADS: + print(f"time: {new_us:.3f} us") + else: + print(f"time H={heads}: {new_us:.3f} us") print(f"flydsl baseline: {old_us:.3f} us") print(f"{backend}: {triton_us:.3f} us") print(f"vs flydsl: {old_us / new_us:.3f}x") print(f"vs {backend}: {triton_us / new_us:.3f}x") -def _profile(): - inp, kv_cache, out = _inputs() +def _profile(heads): + inp, kv_cache, out = _inputs(heads) with torch.inference_mode(): _launch(inp, kv_cache, out) torch.cuda.synchronize() - print("profile launch: pass") + print(f"profile launch H={heads}: pass") def main(): parser = argparse.ArgumentParser() parser.add_argument("--profile", action="store_true") parser.add_argument("--time", action="store_true") + parser.add_argument( + "--heads", + type=int, + choices=HEADS, + default=None, + help="Restrict to one head count (default: both; --time still prints H=64 first)", + ) args = parser.parse_args() - if args.profile: - _profile() - elif args.time: - _benchmark() + heads = (args.heads,) if args.heads is not None else HEADS + if args.time: + # H=64 first so WaveScope still sees `time: {us} us` as the production line. + timed = tuple(h for h in (DEFAULT_HEADS,) + heads if h in heads) + seen = set() + ordered = [] + for h in timed: + if h not in seen: + ordered.append(h) + seen.add(h) + for h in ordered: + _benchmark(h) + elif args.profile: + for h in heads: + _profile(h) else: - test_gfx950_indexer_mapping() + for h in heads: + test_gfx950_indexer_mapping(h) + test_gfx950_nq1_and_short_pages(h) if __name__ == "__main__": From 8adbf328ca51da9cc94ec22ef15daaf854d03c65 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Mon, 7 Sep 2026 14:15:56 +0000 Subject: [PATCH 32/52] feat(flydsl): compile-time next_n for gfx950 paged MQA logits Drop the unused second query row on next_n=1 so MFMA/VALU and store waits match the actual decode shape. Co-authored-by: Cursor --- .../_fp8_paged_mqa_logits_gfx950.py | 8 +++- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 45 +++++++++---------- 2 files changed, 28 insertions(+), 25 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py index 5a4207b2b38..076db967a03 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py @@ -23,10 +23,14 @@ B_RING = KV_BLOCK_SIZE // MFMA_N # Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). PAGE_VMEM_LOADS = B_RING * 3 -# Predicated logit stores: 2 query rows × 4 tiles. Hot-loop pages are full, so -# these 8 VM ops land after the next-page loads and can stay in flight. +# Predicated logit stores: next_n query rows × 4 tiles. Hot-loop pages are full, +# so these VM ops land after the next-page loads and can stay in flight. PAGE_VMEM_STORES = NEXT_N_MAX * B_RING + +def page_vmem_stores(next_n): + return int(next_n) * B_RING + _NEUTRAL_E8M0 = 0x7F7F7F7F diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 51469404fd4..1ed745095b5 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -22,10 +22,8 @@ HEAD_DIM, MFMA_M, MFMA_N, - NEXT_N_MAX, KV_BLOCK_SIZE, PAGE_VMEM_LOADS, - PAGE_VMEM_STORES, SUPPORTED_HEADS, guarded_store, imin, @@ -36,6 +34,7 @@ reduce_scores, schedule_mfma_valu_pairs, uceildiv, + page_vmem_stores, wait_vmcnt, ) from ._mqa_logits_common import ( @@ -52,11 +51,13 @@ def _unwrap(value): return value.ir_value() if hasattr(value, "ir_value") else value -def _build_kernel(*, index_dim: int, num_heads: int): +def _build_kernel(*, index_dim: int, num_heads: int, next_n: int): m_tiles = num_heads // MFMA_M + n_rows = int(next_n) + store_vmem = page_vmem_stores(n_rows) kernel_name = ( f"fp8_paged_mqa_logits_gfx950_H{num_heads}_D{HEAD_DIM}_" - f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{NEXT_N_MAX}_w1_nt_ps_flydsl" + f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{n_rows}_w1_nt_ps_flydsl" ) @flyc.kernel(name=kernel_name, known_block_size=[64, 1, 1]) @@ -67,7 +68,6 @@ def kernel( out_logits: fx.Tensor, context_lens: fx.Tensor, kv_indices: fx.Tensor, - next_n: fx.Int32, batch_size: fx.Int32, split_kv: fx.Int32, stride_q_batch: fx.Int32, @@ -98,18 +98,19 @@ def kernel( col_lo = page_lo * KV_BLOCK_SIZE col_hi = page_hi * KV_BLOCK_SIZE + next_n_c = fx.Int32(n_rows) a_rows = [ [None for _ in range_constexpr(m_tiles)] - for _ in range_constexpr(NEXT_N_MAX) + for _ in range_constexpr(n_rows) ] w_rows = [ [[None for _ in range_constexpr(DREG)] for _ in range_constexpr(m_tiles)] - for _ in range_constexpr(NEXT_N_MAX) + for _ in range_constexpr(n_rows) ] - for row in range_constexpr(NEXT_N_MAX): - q_row = imin(fx.Int32(row), next_n - 1) + for row in range_constexpr(n_rows): + q_row = fx.Int32(row) q_base = pid_batch * stride_q_batch + q_row * stride_q_next_n - out_row = pid_batch * next_n + q_row + out_row = pid_batch * next_n_c + q_row for mi in range_constexpr(m_tiles): h = mi * MFMA_M + lane_mod_16 byte_base = q_base + h * stride_q_heads @@ -162,12 +163,11 @@ def _unpack_state(state): def _store_row(row, col, scores, scale): q_row = fx.Int32(row) - q_limit = context_len - next_n + q_row - out_row = pid_batch * next_n + q_row + q_limit = context_len - next_n_c + q_row + out_row = pid_batch * next_n_c + q_row logit = reduce_scores(scores, w_rows[row], scale, m_tiles=m_tiles) writer = ( (lane_div_16 == 0) - & (q_row < next_n) & (col < context_len) & (col <= q_limit) ) @@ -185,17 +185,17 @@ def _compute_page(page_col, b_slots, scale_slots): col = page_col + slot * MFMA_N + lane_mod_16 scores = [ mfma_scores(a_rows[row], b_slots[slot], m_tiles=m_tiles) - for row in range_constexpr(NEXT_N_MAX) + for row in range_constexpr(n_rows) ] if slot > 0: schedule_mfma_valu_pairs(m_tiles=m_tiles) - for row in range_constexpr(NEXT_N_MAX): + for row in range_constexpr(n_rows): _store_row(row, prev_col, prev_scores[row], prev_scale) prev_scores = scores prev_col = col prev_scale = scale_slots[slot] schedule_mfma_valu_pairs(m_tiles=m_tiles) - for row in range_constexpr(NEXT_N_MAX): + for row in range_constexpr(n_rows): _store_row(row, prev_col, prev_scores[row], prev_scale) rocdl.sched_barrier(0) @@ -215,7 +215,7 @@ def _compute_page(page_col, b_slots, scale_slots): # the current page (or leftover stores) has retired. wait_vmcnt(PAGE_VMEM_LOADS) _compute_page(fx.Int32(page_col), current_b, current_scale) - wait_vmcnt(PAGE_VMEM_STORES) + wait_vmcnt(store_vmem) results = yield _pack_state(next_b, next_scale) final_b, final_scale = _unpack_state(results) _compute_page(stop, final_b, final_scale) @@ -232,7 +232,6 @@ def launch( context_lens, kv_indices, grid_blocks, - next_n, batch_size, split_kv, stride_q_batch, @@ -250,7 +249,6 @@ def launch( out_logits, context_lens, kv_indices, - next_n, batch_size, split_kv, stride_q_batch, @@ -270,8 +268,8 @@ def launch( @lru_cache(maxsize=8) -def _compile(*, index_dim: int, num_heads: int): - return _build_kernel(index_dim=index_dim, num_heads=num_heads) +def _compile(*, index_dim: int, num_heads: int, next_n: int): + return _build_kernel(index_dim=index_dim, num_heads=num_heads, next_n=next_n) def flydsl_fp8_paged_mqa_logits_gfx950( @@ -329,7 +327,9 @@ def flydsl_fp8_paged_mqa_logits_gfx950( else max(1, min(max_pages, int(SplitKV))) ) grid_blocks = batch_size * split_kv - launcher = _compile(index_dim=int(index_dim), num_heads=int(heads)) + launcher = _compile( + index_dim=int(index_dim), num_heads=int(heads), next_n=int(next_n) + ) launcher.compile_hints = {**DEFAULT_COMPILE_HINTS, "waves_per_eu": 2} stream = stream or torch.cuda.current_stream(q_fp8.device) with torch.cuda.device(q_fp8.device.index): @@ -342,7 +342,6 @@ def flydsl_fp8_paged_mqa_logits_gfx950( context_lens, kv_indices, int(grid_blocks), - int(next_n), int(batch_size), int(split_kv), int(q_fp8.stride(0)), From 6c1c496cdb85ea4e1b6efaa742e68e439a6c56a0 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Mon, 7 Sep 2026 14:40:33 +0000 Subject: [PATCH 33/52] feat(flydsl): dual-bank page pipeline for gfx950 paged MQA logits Keep next-page K in fixed register banks instead of scf.for yield copies, store one wave-wide page of logits, and prefetch the page table so gather latency stays overlapped. Co-authored-by: Cursor --- .../_fp8_paged_mqa_logits_gfx950.py | 9 +- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 162 +++++++++++++----- 2 files changed, 124 insertions(+), 47 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py index 076db967a03..29110448fe8 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py @@ -15,7 +15,6 @@ SUPPORTED_HEADS = (32, 64) HEAD_DIM = 128 KV_BLOCK_SIZE = 64 -NEXT_N_MAX = 2 MFMA_M = 16 MFMA_N = 16 MFMA_K = 128 @@ -23,13 +22,9 @@ B_RING = KV_BLOCK_SIZE // MFMA_N # Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). PAGE_VMEM_LOADS = B_RING * 3 -# Predicated logit stores: next_n query rows × 4 tiles. Hot-loop pages are full, -# so these VM ops land after the next-page loads and can stay in flight. -PAGE_VMEM_STORES = NEXT_N_MAX * B_RING - - def page_vmem_stores(next_n): - return int(next_n) * B_RING + """One wave-wide contiguous store per query row and physical page.""" + return int(next_n) _NEUTRAL_E8M0 = 0x7F7F7F7F diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 1ed745095b5..e25cd2c6385 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -57,7 +57,7 @@ def _build_kernel(*, index_dim: int, num_heads: int, next_n: int): store_vmem = page_vmem_stores(n_rows) kernel_name = ( f"fp8_paged_mqa_logits_gfx950_H{num_heads}_D{HEAD_DIM}_" - f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{n_rows}_w1_nt_ps_flydsl" + f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{n_rows}_db2_ws1_pfp_sb0_nt_ps_flydsl" ) @flyc.kernel(name=kernel_name, known_block_size=[64, 1, 1]) @@ -126,13 +126,15 @@ def kernel( for ii in range_constexpr(DREG): w_rows[row][mi][ii] = fx.Float32(weight_vec[ii]) - def _issue_page(page_col): + def _load_physical(page_col): table_idx = pid_batch * max_block_len + udiv(page_col, fx.Int32(KV_BLOCK_SIZE)) - physical = fx.Int32( + return fx.Int32( buffer_ops.buffer_load( table_t.rsrc, table_idx, vec_width=1, dtype=T.i32, is_scalar=True ) ) + + def _issue_physical(physical): b_slots, scale_slots = [], [] for slot in range_constexpr(B_RING): token = slot * MFMA_N + lane_mod_16 @@ -151,38 +153,82 @@ def _issue_page(page_col): ) return b_slots, scale_slots - def _pack_state(b_slots, scale_slots): - return [_unwrap(v) for v in b_slots + scale_slots] + def _issue_page(page_col): + return _issue_physical(_load_physical(page_col)) - def _unpack_state(state): - b_slots = [fx.Vector(state[i]) for i in range_constexpr(B_RING)] - scale_slots = [ - fx.Float32(state[B_RING + i]) for i in range_constexpr(B_RING) - ] - return b_slots, scale_slots + def _make_page_bank(): + return ( + [ + fx.make_rmem_tensor(DREG * 2, fx.Int32) + for _ in range_constexpr(B_RING) + ], + [ + fx.make_rmem_tensor(1, fx.Float32) + for _ in range_constexpr(B_RING) + ], + ) + + def _store_page_bank(bank, b_slots, scale_slots): + b_bank, scale_bank = bank + for slot in range_constexpr(B_RING): + b_bank[slot].store(b_slots[slot]) + scale_bank[slot].store( + fx.Vector.from_elements( + [_unwrap(scale_slots[slot])], fx.Float32 + ) + ) + + def _load_page_bank(bank): + b_bank, scale_bank = bank + return ( + [ + fx.Vector(b_bank[slot].load()) + for slot in range_constexpr(B_RING) + ], + [ + fx.Float32(fx.Vector(scale_bank[slot].load())[0]) + for slot in range_constexpr(B_RING) + ], + ) + + def _make_physical_bank(): + return fx.make_rmem_tensor(1, fx.Int32) - def _store_row(row, col, scores, scale): + def _store_physical(bank, physical): + bank.store(fx.Vector.from_elements([_unwrap(physical)], fx.Int32)) + + def _load_physical_bank(bank): + return fx.Int32(fx.Vector(bank.load())[0]) + + def _reduce_row(row, scores, scale): + return reduce_scores(scores, w_rows[row], scale, m_tiles=m_tiles) + + def _store_page_row(row, page_col, logits): q_row = fx.Int32(row) q_limit = context_len - next_n_c + q_row out_row = pid_batch * next_n_c + q_row - logit = reduce_scores(scores, w_rows[row], scale, m_tiles=m_tiles) + col = page_col + lane + logit = logits[0] + for slot in range_constexpr(1, B_RING): + logit = (lane_div_16 == slot).select(logits[slot], logit) writer = ( - (lane_div_16 == 0) - & (col < context_len) + (col < context_len) & (col <= q_limit) ) + def _write(_row=out_row, _col=col, _value=logit): out_t[_row * stride_out + _col] = _value guarded_store(writer, _write) def _compute_page(page_col, b_slots, scale_slots): - rocdl.sched_barrier(0) prev_scores = None - prev_col = None prev_scale = None + page_logits = [ + [None for _ in range_constexpr(B_RING)] + for _ in range_constexpr(n_rows) + ] for slot in range_constexpr(B_RING): - col = page_col + slot * MFMA_N + lane_mod_16 scores = [ mfma_scores(a_rows[row], b_slots[slot], m_tiles=m_tiles) for row in range_constexpr(n_rows) @@ -190,38 +236,74 @@ def _compute_page(page_col, b_slots, scale_slots): if slot > 0: schedule_mfma_valu_pairs(m_tiles=m_tiles) for row in range_constexpr(n_rows): - _store_row(row, prev_col, prev_scores[row], prev_scale) + page_logits[row][slot - 1] = _reduce_row( + row, prev_scores[row], prev_scale + ) prev_scores = scores - prev_col = col prev_scale = scale_slots[slot] schedule_mfma_valu_pairs(m_tiles=m_tiles) for row in range_constexpr(n_rows): - _store_row(row, prev_col, prev_scores[row], prev_scale) - rocdl.sched_barrier(0) + page_logits[row][B_RING - 1] = _reduce_row( + row, prev_scores[row], prev_scale + ) + _store_page_row(row, page_col, page_logits[row]) if page_lo < page_hi: + current_bank = _make_page_bank() + next_bank = _make_page_bank() b_init, scale_init = _issue_page(col_lo) - if page_lo + 1 < page_hi: - stop = col_hi - KV_BLOCK_SIZE - init_state = _pack_state(b_init, scale_init) - for page_col, state in range( - col_lo, stop, fx.Int32(KV_BLOCK_SIZE), init=init_state - ): - current_b, current_scale = _unpack_state(state) - next_b, next_scale = _issue_page( - fx.Int32(page_col) + fx.Int32(KV_BLOCK_SIZE) + _store_page_bank(current_bank, b_init, scale_init) + current_physical = _make_physical_bank() + next_physical = _make_physical_bank() + if col_lo + fx.Int32(KV_BLOCK_SIZE) < col_hi: + _store_physical( + next_physical, + _load_physical(col_lo + fx.Int32(KV_BLOCK_SIZE)), + ) + page_col = col_lo + while page_col + fx.Int32(KV_BLOCK_SIZE) < col_hi: + next_b, next_scale = _issue_physical( + _load_physical_bank(next_physical) + ) + _store_page_bank(next_bank, next_b, next_scale) + has_following = page_col + fx.Int32(2 * KV_BLOCK_SIZE) < col_hi + if has_following: + _store_physical( + current_physical, + _load_physical(page_col + fx.Int32(2 * KV_BLOCK_SIZE)), + ) + # Leave the next-page gather outstanding while consuming current. + wait_vmcnt(PAGE_VMEM_LOADS) + current_b, current_scale = _load_page_bank(current_bank) + _compute_page(page_col, current_b, current_scale) + + if has_following: + following_b, following_scale = _issue_physical( + _load_physical_bank(current_physical) ) - # 12 next-page loads in flight: wait until those 12 remain so - # the current page (or leftover stores) has retired. - wait_vmcnt(PAGE_VMEM_LOADS) - _compute_page(fx.Int32(page_col), current_b, current_scale) + _store_page_bank(current_bank, following_b, following_scale) + if page_col + fx.Int32(3 * KV_BLOCK_SIZE) < col_hi: + _store_physical( + next_physical, + _load_physical( + page_col + fx.Int32(3 * KV_BLOCK_SIZE) + ), + ) + # following-page loads plus current-page stores can remain. + wait_vmcnt(PAGE_VMEM_LOADS + store_vmem) + else: + # Only current-page stores may remain. wait_vmcnt(store_vmem) - results = yield _pack_state(next_b, next_scale) - final_b, final_scale = _unpack_state(results) - _compute_page(stop, final_b, final_scale) - else: + next_b, next_scale = _load_page_bank(next_bank) + _compute_page( + page_col + fx.Int32(KV_BLOCK_SIZE), next_b, next_scale + ) + page_col = page_col + fx.Int32(2 * KV_BLOCK_SIZE) + + if page_col < col_hi: wait_vmcnt(0) - _compute_page(col_lo, b_init, scale_init) + current_b, current_scale = _load_page_bank(current_bank) + _compute_page(page_col, current_b, current_scale) @flyc.jit def launch( From c1ee3ec73fdc31d097d43cbd7088bf22f292e657 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Mon, 7 Sep 2026 14:55:44 +0000 Subject: [PATCH 34/52] feat(flydsl): dispatch gfx950 paged MQA logits as the default decode path Route H32/H64 D128 KVB64 preshuffle next_n in {1,2} through the 16x16x128 kernel, keep the 32x32x64 implementation as the generic fallback, and satisfy Black/Ruff. Co-authored-by: Cursor --- aiter/ops/flydsl/__init__.py | 5 + .../_fp8_paged_mqa_logits_gfx950.py | 10 +- .../mqa_logits/fp8_paged_mqa_logits.py | 120 +++++++++++++++++- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 55 +++----- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 11 +- 5 files changed, 156 insertions(+), 45 deletions(-) diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 5bca1889732..93417478fc7 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -54,6 +54,10 @@ ".kernels.mqa_logits.fp8_paged_mqa_logits", "flydsl_fp8_paged_mqa_logits", ), + "flydsl_fp8_paged_mqa_logits_generic": ( + ".kernels.mqa_logits.fp8_paged_mqa_logits", + "flydsl_fp8_paged_mqa_logits_generic", + ), "flydsl_fp8_paged_mqa_logits_gfx950": ( ".kernels.mqa_logits.fp8_paged_mqa_logits_gfx950", "flydsl_fp8_paged_mqa_logits_gfx950", @@ -94,6 +98,7 @@ "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", "flydsl_fp8_paged_mqa_logits", + "flydsl_fp8_paged_mqa_logits_generic", "flydsl_fp8_paged_mqa_logits_gfx950", "flydsl_hgemm", "flydsl_mla_reduce_v1", diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py index 29110448fe8..b4309c4a11b 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py @@ -22,10 +22,13 @@ B_RING = KV_BLOCK_SIZE // MFMA_N # Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). PAGE_VMEM_LOADS = B_RING * 3 + + def page_vmem_stores(next_n): """One wave-wide contiguous store per query row and physical page.""" return int(next_n) + _NEUTRAL_E8M0 = 0x7F7F7F7F @@ -100,8 +103,11 @@ def load_preshuffled_k_pack( kv_i32.rsrc, base, vec_width=4, dtype=T.i32, cache_modifier=2 ) hi = buffer_ops.buffer_load( - kv_i32.rsrc, base + fx.Int32(MFMA_N * 16 // 4), - vec_width=4, dtype=T.i32, cache_modifier=2, + kv_i32.rsrc, + base + fx.Int32(MFMA_N * 16 // 4), + vec_width=4, + dtype=T.i32, + cache_modifier=2, ) return _concat_i32x4(lo, hi) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 95bf6326c70..e3e92aab497 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -19,6 +19,12 @@ are interchangeable. ``Preshuffle=True`` consumes the production ``shuffle_weight(layout=(16,16))`` KV layout (needs ``KVBlockSize % 16 == 0``). +The public ``flydsl_fp8_paged_mqa_logits`` entry dispatches the gfx950 +16x16x128 mapping (``fp8_paged_mqa_logits_gfx950``) for the production decode +shape ``H in {32,64}``, ``D=128``, ``KVBlockSize=64``, ``next_n in {1,2}``, +preshuffled cache. Other shapes keep this 32x32x64 generic kernel. Force the +generic path with ``FLYDSL_FP8_PAGED_MQA_LOGITS_GENERIC=1``. + Cache layout, per physical block of ``KVBlockSize`` tokens: the fp8 key rows (``D`` bytes each) grouped first, then the f32 dequant scales. Each lane resolves its own column through ``kv_indices[b, p // KVBlockSize]``, so a column @@ -404,7 +410,46 @@ def compile_fp8_paged_mqa_logits( return launcher -def flydsl_fp8_paged_mqa_logits( +def _use_gfx950_fastpath( + q_fp8, + kv_cache, + *, + Preshuffle, + KVBlockSize, + ChunkK, + WavePerEU, + variant, +): + """True for the production decode shape the 16x16x128 gfx950 kernel covers.""" + if os.environ.get("FLYDSL_FP8_PAGED_MQA_LOGITS_GENERIC", "").strip() in ( + "1", + "true", + "True", + ): + return False + env_variant = os.environ.get("FLYDSL_FP8_PAGED_MQA_LOGITS_VARIANT") + if variant not in (None, DEFAULT_VARIANT) or ( + env_variant and env_variant != DEFAULT_VARIANT + ): + return False + if not Preshuffle or int(KVBlockSize) != 64: + return False + if int(ChunkK) != _BLOCK_KV or int(WavePerEU) != 2: + return False + if get_gfx() != "gfx950": + return False + _, next_n, num_heads, head_size = q_fp8.shape + if num_heads not in (32, 64) or head_size != 128 or next_n not in (1, 2): + return False + from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype + + if q_fp8.dtype != get_fp8_e4m3_dtype(): + return False + _, block_size, one, index_dim = kv_cache.shape + return block_size == 64 and one == 1 and index_dim == 132 + + +def flydsl_fp8_paged_mqa_logits_generic( q_fp8, kv_cache, weights, @@ -422,9 +467,11 @@ def flydsl_fp8_paged_mqa_logits( variant=None, stream=None, ): - """FlyDSL paged FP8 MQA logits (decode) -- KVBlockSize>=1 with SplitKV. + """Generic 32x32x64 FlyDSL paged FP8 MQA logits (decode). Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. + Production H32/H64 D128 KVB64 preshuffle decode uses the gfx950 mapping via + ``flydsl_fp8_paged_mqa_logits`` instead. q_fp8: [batch, next_n, heads, hidden_dim], float8 e4m3fn (gfx950 native) kv_cache: [num_blocks, KVBlockSize, 1, index_dim] uint8, co-packed per @@ -500,7 +547,7 @@ def flydsl_fp8_paged_mqa_logits( ), f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" assert ( get_gfx() == "gfx950" - ), f"flydsl_fp8_paged_mqa_logits targets gfx950 (32x32x64 MFMA); got {get_gfx()}" + ), f"flydsl_fp8_paged_mqa_logits_generic targets gfx950 (32x32x64 MFMA); got {get_gfx()}" variant = _resolve_variant(variant) @@ -558,3 +605,70 @@ def flydsl_fp8_paged_mqa_logits( ) return out_logits + + +def flydsl_fp8_paged_mqa_logits( + q_fp8, + kv_cache, + weights, + out_logits, + context_lens, + kv_indices, + max_model_len, + *, + Preshuffle=False, + KVBlockSize=1, + ChunkK=_BLOCK_KV, + SplitKV=None, + WavePerEU=2, + TotalCuCount=None, + variant=None, + stream=None, +): + """FlyDSL paged FP8 MQA logits (decode) -- KVBlockSize>=1 with SplitKV. + + Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. + Dispatches the gfx950 16x16x128 kernel for H32/H64 D128 KVB64 preshuffle + ``next_n in {1,2}``; otherwise the generic 32x32x64 kernel. + """ + if _use_gfx950_fastpath( + q_fp8, + kv_cache, + Preshuffle=Preshuffle, + KVBlockSize=KVBlockSize, + ChunkK=ChunkK, + WavePerEU=WavePerEU, + variant=variant, + ): + from .fp8_paged_mqa_logits_gfx950 import flydsl_fp8_paged_mqa_logits_gfx950 + + return flydsl_fp8_paged_mqa_logits_gfx950( + q_fp8, + kv_cache, + weights, + out_logits, + context_lens, + kv_indices, + max_model_len, + KVBlockSize=KVBlockSize, + SplitKV=SplitKV, + TotalCuCount=TotalCuCount, + stream=stream, + ) + return flydsl_fp8_paged_mqa_logits_generic( + q_fp8, + kv_cache, + weights, + out_logits, + context_lens, + kv_indices, + max_model_len, + Preshuffle=Preshuffle, + KVBlockSize=KVBlockSize, + ChunkK=ChunkK, + SplitKV=SplitKV, + WavePerEU=WavePerEU, + TotalCuCount=TotalCuCount, + variant=variant, + stream=stream, + ) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index e25cd2c6385..30159d0fe4b 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -1,14 +1,18 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx950 H32/H64 D128/KVB64 preshuffled paged FP8 MQA-logits kernel.""" +"""gfx950 H32/H64 D128/KVB64 preshuffled paged FP8 MQA-logits kernel. + +Default decode path for ``flydsl_fp8_paged_mqa_logits`` on this shape. +``flydsl_fp8_paged_mqa_logits_gfx950`` remains the explicit entry. +""" from functools import lru_cache import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl.expr import arith, range_constexpr, rocdl +from flydsl.expr import arith, range_constexpr from flydsl.expr.typing import T from aiter.jit.utils.chip_info import get_gfx @@ -20,9 +24,9 @@ B_RING, DREG, HEAD_DIM, + KV_BLOCK_SIZE, MFMA_M, MFMA_N, - KV_BLOCK_SIZE, PAGE_VMEM_LOADS, SUPPORTED_HEADS, guarded_store, @@ -31,10 +35,10 @@ load_preshuffled_k_pack, load_q_pack, mfma_scores, + page_vmem_stores, reduce_scores, schedule_mfma_valu_pairs, uceildiv, - page_vmem_stores, wait_vmcnt, ) from ._mqa_logits_common import ( @@ -100,8 +104,7 @@ def kernel( next_n_c = fx.Int32(n_rows) a_rows = [ - [None for _ in range_constexpr(m_tiles)] - for _ in range_constexpr(n_rows) + [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(n_rows) ] w_rows = [ [[None for _ in range_constexpr(DREG)] for _ in range_constexpr(m_tiles)] @@ -127,7 +130,9 @@ def kernel( w_rows[row][mi][ii] = fx.Float32(weight_vec[ii]) def _load_physical(page_col): - table_idx = pid_batch * max_block_len + udiv(page_col, fx.Int32(KV_BLOCK_SIZE)) + table_idx = pid_batch * max_block_len + udiv( + page_col, fx.Int32(KV_BLOCK_SIZE) + ) return fx.Int32( buffer_ops.buffer_load( table_t.rsrc, table_idx, vec_width=1, dtype=T.i32, is_scalar=True @@ -162,10 +167,7 @@ def _make_page_bank(): fx.make_rmem_tensor(DREG * 2, fx.Int32) for _ in range_constexpr(B_RING) ], - [ - fx.make_rmem_tensor(1, fx.Float32) - for _ in range_constexpr(B_RING) - ], + [fx.make_rmem_tensor(1, fx.Float32) for _ in range_constexpr(B_RING)], ) def _store_page_bank(bank, b_slots, scale_slots): @@ -173,18 +175,13 @@ def _store_page_bank(bank, b_slots, scale_slots): for slot in range_constexpr(B_RING): b_bank[slot].store(b_slots[slot]) scale_bank[slot].store( - fx.Vector.from_elements( - [_unwrap(scale_slots[slot])], fx.Float32 - ) + fx.Vector.from_elements([_unwrap(scale_slots[slot])], fx.Float32) ) def _load_page_bank(bank): b_bank, scale_bank = bank return ( - [ - fx.Vector(b_bank[slot].load()) - for slot in range_constexpr(B_RING) - ], + [fx.Vector(b_bank[slot].load()) for slot in range_constexpr(B_RING)], [ fx.Float32(fx.Vector(scale_bank[slot].load())[0]) for slot in range_constexpr(B_RING) @@ -211,10 +208,7 @@ def _store_page_row(row, page_col, logits): logit = logits[0] for slot in range_constexpr(1, B_RING): logit = (lane_div_16 == slot).select(logits[slot], logit) - writer = ( - (col < context_len) - & (col <= q_limit) - ) + writer = (col < context_len) & (col <= q_limit) def _write(_row=out_row, _col=col, _value=logit): out_t[_row * stride_out + _col] = _value @@ -225,8 +219,7 @@ def _compute_page(page_col, b_slots, scale_slots): prev_scores = None prev_scale = None page_logits = [ - [None for _ in range_constexpr(B_RING)] - for _ in range_constexpr(n_rows) + [None for _ in range_constexpr(B_RING)] for _ in range_constexpr(n_rows) ] for slot in range_constexpr(B_RING): scores = [ @@ -262,9 +255,7 @@ def _compute_page(page_col, b_slots, scale_slots): ) page_col = col_lo while page_col + fx.Int32(KV_BLOCK_SIZE) < col_hi: - next_b, next_scale = _issue_physical( - _load_physical_bank(next_physical) - ) + next_b, next_scale = _issue_physical(_load_physical_bank(next_physical)) _store_page_bank(next_bank, next_b, next_scale) has_following = page_col + fx.Int32(2 * KV_BLOCK_SIZE) < col_hi if has_following: @@ -285,9 +276,7 @@ def _compute_page(page_col, b_slots, scale_slots): if page_col + fx.Int32(3 * KV_BLOCK_SIZE) < col_hi: _store_physical( next_physical, - _load_physical( - page_col + fx.Int32(3 * KV_BLOCK_SIZE) - ), + _load_physical(page_col + fx.Int32(3 * KV_BLOCK_SIZE)), ) # following-page loads plus current-page stores can remain. wait_vmcnt(PAGE_VMEM_LOADS + store_vmem) @@ -295,9 +284,7 @@ def _compute_page(page_col, b_slots, scale_slots): # Only current-page stores may remain. wait_vmcnt(store_vmem) next_b, next_scale = _load_page_bank(next_bank) - _compute_page( - page_col + fx.Int32(KV_BLOCK_SIZE), next_b, next_scale - ) + _compute_page(page_col + fx.Int32(KV_BLOCK_SIZE), next_b, next_scale) page_col = page_col + fx.Int32(2 * KV_BLOCK_SIZE) if page_col < col_hi: @@ -383,7 +370,7 @@ def flydsl_fp8_paged_mqa_logits_gfx950( raise ValueError(f"q_fp8 must be native FP8 E4M3, got {q_fp8.dtype}") if kv_cache.dtype != torch.uint8: raise ValueError("kv_cache must contain preshuffled uint8 FP8 data") - num_blocks, block_size, one, index_dim = kv_cache.shape + _, block_size, one, index_dim = kv_cache.shape if block_size != KV_BLOCK_SIZE or one != 1 or index_dim != HEAD_DIM + 4: raise ValueError(f"unexpected KV cache shape {tuple(kv_cache.shape)}") diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index e86612b5a27..b76f05ec4e2 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -10,7 +10,7 @@ from aiter.jit.utils.chip_info import get_gfx from aiter.ops.flydsl import ( flydsl_fp8_paged_mqa_logits, - flydsl_fp8_paged_mqa_logits_gfx950, + flydsl_fp8_paged_mqa_logits_generic, ) from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.test_common import run_perftest @@ -42,14 +42,12 @@ def _inputs(heads, next_n=NEXT_N, kv_len=KV_LEN, batch=BATCH): get_fp8_e4m3_dtype(), block_size=KV_BLOCK_SIZE, ) - kv_cache, out = _kernel_inputs( - inp, batch, next_n, HEAD_DIM, True, KV_BLOCK_SIZE - ) + kv_cache, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM, True, KV_BLOCK_SIZE) return inp, kv_cache, out def _launch(inp, kv_cache, out): - return flydsl_fp8_paged_mqa_logits_gfx950( + return flydsl_fp8_paged_mqa_logits( inp.q_fp8, kv_cache, inp.weights, @@ -57,6 +55,7 @@ def _launch(inp, kv_cache, out): inp.context_lens, inp.block_tables, inp.max_model_len, + Preshuffle=True, KVBlockSize=KV_BLOCK_SIZE, ) @@ -113,7 +112,7 @@ def new_kernel(): return _launch(inp, kv_cache, out_new) def old_kernel(): - return flydsl_fp8_paged_mqa_logits( + return flydsl_fp8_paged_mqa_logits_generic( inp.q_fp8, kv_cache, inp.weights, From 9661da3a7f6a6aaa99f0ae6907cb84e5d9d77a90 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 8 Sep 2026 12:08:56 +0000 Subject: [PATCH 35/52] refactor(flydsl): make gfx950 the only paged FP8 MQA logits path Drop the generic 32x32x64 kernel and fold helpers into the gfx950 decoder so H32/H64 stays on the tuned 16x16x128 pipeline, with i64 output-row addressing and idiomatic FlyDSL MMA/scheduling that still matches the previous ISA. Co-authored-by: Cursor --- aiter/ops/flydsl/__init__.py | 17 +- .../_fp8_paged_mqa_logits_gfx950.py | 158 ---- .../kernels/mqa_logits/_mqa_logits_common.py | 181 ----- .../mqa_logits/fp8_paged_mqa_logits.py | 673 +----------------- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 291 ++++++-- .../test_flydsl_fp8_paged_mqa_logits.py | 148 ++-- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 66 +- 7 files changed, 323 insertions(+), 1211 deletions(-) delete mode 100644 aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py delete mode 100644 aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index 93417478fc7..558e1ace687 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -33,14 +33,6 @@ ".kernels.mqa_logits.fp8_mqa_logits", "KERNEL_VARIANTS", ), - "FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT": ( - ".kernels.mqa_logits.fp8_paged_mqa_logits", - "DEFAULT_VARIANT", - ), - "FP8_PAGED_MQA_LOGITS_VARIANTS": ( - ".kernels.mqa_logits.fp8_paged_mqa_logits", - "KERNEL_VARIANTS", - ), "compute_varqlen_windows": ( ".kernels.mqa_logits.pa_mqa_logits_fp4_prefill", "compute_varqlen_windows", @@ -51,13 +43,9 @@ "flydsl_fp8_mqa_logits", ), "flydsl_fp8_paged_mqa_logits": ( - ".kernels.mqa_logits.fp8_paged_mqa_logits", + ".kernels.mqa_logits.fp8_paged_mqa_logits_gfx950", "flydsl_fp8_paged_mqa_logits", ), - "flydsl_fp8_paged_mqa_logits_generic": ( - ".kernels.mqa_logits.fp8_paged_mqa_logits", - "flydsl_fp8_paged_mqa_logits_generic", - ), "flydsl_fp8_paged_mqa_logits_gfx950": ( ".kernels.mqa_logits.fp8_paged_mqa_logits_gfx950", "flydsl_fp8_paged_mqa_logits_gfx950", @@ -91,14 +79,11 @@ __all__ = [ "FP8_MQA_LOGITS_DEFAULT_VARIANT", "FP8_MQA_LOGITS_VARIANTS", - "FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT", - "FP8_PAGED_MQA_LOGITS_VARIANTS", "GateMode", "compute_varqlen_windows", "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", "flydsl_fp8_paged_mqa_logits", - "flydsl_fp8_paged_mqa_logits_generic", "flydsl_fp8_paged_mqa_logits_gfx950", "flydsl_hgemm", "flydsl_mla_reduce_v1", diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py deleted file mode 100644 index b4309c4a11b..00000000000 --- a/aiter/ops/flydsl/kernels/mqa_logits/_fp8_paged_mqa_logits_gfx950.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""16x16x128 FP8 primitives for the gfx950 paged-MQA indexer fast path.""" - -import flydsl.compiler as flyc -import flydsl.expr as fx -from flydsl.expr import arith, range_constexpr, rocdl -from flydsl.expr.typing import T - -from .. import buffer_ops - -Vec = fx.Vector - -SUPPORTED_HEADS = (32, 64) -HEAD_DIM = 128 -KV_BLOCK_SIZE = 64 -MFMA_M = 16 -MFMA_N = 16 -MFMA_K = 128 -DREG = 4 -B_RING = KV_BLOCK_SIZE // MFMA_N -# Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). -PAGE_VMEM_LOADS = B_RING * 3 - - -def page_vmem_stores(next_n): - """One wave-wide contiguous store per query row and physical page.""" - return int(next_n) - - -_NEUTRAL_E8M0 = 0x7F7F7F7F - - -def wait_vmcnt(n): - """Explicit s_waitcnt vmcnt(n). Leaves n VM ops outstanding.""" - rocdl.sched_barrier(0) - rocdl.s_waitcnt(vmcnt=int(n)) - rocdl.sched_barrier(0) - - -def uceildiv(a, b): - a, b = fx.Int32(a), fx.Int32(b) - return fx.Int32((fx.Uint32(a) + fx.Uint32(b) - 1) // fx.Uint32(b)) - - -def imin(a, b): - a, b = fx.Int32(a), fx.Int32(b) - return (a <= b).select(a, b) - - -def guarded_store(pred, store_fn): - """Predicated side effect (FlyDSL ``@flyc.jit`` branch on a runtime mask).""" - - @flyc.jit - def _guarded(_pred=pred, _store=store_fn): - if _pred: - _store() - - _guarded() - - -def _concat_i32x4(lo, hi): - lo, hi = Vec(lo), Vec(hi) - return Vec.from_elements( - [Vec(lo)[i].ir_value() for i in range_constexpr(4)] - + [Vec(hi)[i].ir_value() for i in range_constexpr(4)], - fx.Int32, - ) - - -def load_q_pack(q_i32, byte_base, lane_div_16): - """Load one row's lane-owned contiguous K32 segment as a 256-bit operand.""" - off = (byte_base + lane_div_16 * 32) // fx.Int32(4) - lo = buffer_ops.buffer_load(q_i32.rsrc, off, vec_width=4, dtype=T.i32) - hi = buffer_ops.buffer_load(q_i32.rsrc, off + 4, vec_width=4, dtype=T.i32) - return _concat_i32x4(lo, hi) - - -def load_preshuffled_k_pack( - kv_i32, - physical, - tile_in_page, - lane_mod_16, - lane_div_16, - *, - index_dim, -): - """Load a shuffle_weight(16,16) K column for 16x16x128 scaled MFMA. - - shuffle_weight stores ``[K/32, K-half, token16, byte16]``. ``lane_div_16`` - selects the lane's K32 segment and the two 16-byte halves form i32x8. - """ - block_i32 = fx.Int32(KV_BLOCK_SIZE * index_dim // 4) - lane_k_bytes = MFMA_K // (64 // MFMA_N) - base = ( - physical * block_i32 - + fx.Int32(tile_in_page * HEAD_DIM * MFMA_N // 4) - + lane_div_16 * fx.Int32(lane_k_bytes * MFMA_N // 4) - + lane_mod_16 * 4 - ) - lo = buffer_ops.buffer_load( - kv_i32.rsrc, base, vec_width=4, dtype=T.i32, cache_modifier=2 - ) - hi = buffer_ops.buffer_load( - kv_i32.rsrc, - base + fx.Int32(MFMA_N * 16 // 4), - vec_width=4, - dtype=T.i32, - cache_modifier=2, - ) - return _concat_i32x4(lo, hi) - - -def load_kv_scale(kv_i32, physical, token_in_page, *, index_dim): - block_i32 = fx.Int32(KV_BLOCK_SIZE * index_dim // 4) - off = physical * block_i32 + fx.Int32(KV_BLOCK_SIZE * HEAD_DIM // 4) + token_in_page - return fx.Float32( - buffer_ops.buffer_load( - kv_i32.rsrc, off, vec_width=1, dtype=T.f32, cache_modifier=2 - ) - ) - - -def mfma_scores(a_tiles, b_pack, *, m_tiles): - """Return one 4-f32 score fragment per 16-head M tile.""" - result_type = Vec.make_type(DREG, fx.Float32) - neutral = arith.constant(_NEUTRAL_E8M0, type=T.i32) - scores = [] - for mi in range_constexpr(m_tiles): - acc = Vec.filled(DREG, 0.0, fx.Float32) - acc = rocdl.mfma_scale_f32_16x16x128_f8f6f4( - result_type, - [a_tiles[mi], b_pack, acc, 0, 0, 0, neutral, 0, neutral], - ) - scores.append(acc) - return scores - - -def reduce_scores(scores, weights, kv_scale, *, m_tiles): - """ReLU, weighted H reduction, positive KV scale, then wave reduction.""" - zero = fx.Float32(0.0) - total = zero - for mi in range_constexpr(m_tiles): - frag = Vec(scores[mi]) - for ii in range_constexpr(DREG): - total = total + fx.Float32(frag[ii]).maximumf(zero) * weights[mi][ii] - total = total * kv_scale - total = total + total.shuffle_xor(16, 64) - total = total + total.shuffle_xor(32, 64) - return total - - -def schedule_mfma_valu_pairs(*, m_tiles): - """Pair each 32-cycle MFMA with the prior tile's 12-op VALU fragment.""" - for _ in range_constexpr(m_tiles): - rocdl.sched_group_barrier(0x008, 1, 0) - rocdl.sched_group_barrier(0x002, DREG * 3, 0) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py b/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py deleted file mode 100644 index d8eff26e26d..00000000000 --- a/aiter/ops/flydsl/kernels/mqa_logits/_mqa_logits_common.py +++ /dev/null @@ -1,181 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""FlyDSL building blocks for the paged FP8 MQA-logits kernel (gfx950 / CDNA4). - -logits[row, n] = sum_h ReLU( * kv_scale[n]) * weights[row, h] -""" - -from functools import lru_cache - -import flydsl.expr as fx -import torch -from flydsl.expr import arith, range_constexpr, rocdl -from flydsl.expr.typing import T - -from ..tensor_shim import GTensor, _to_raw - -Vec = fx.Vector - -# 32x32x64 scaled fp8 MFMA (gfx950 / CDNA4). A/B are vector<8xi32> per lane -# (32 fp8 bytes), C/D are 16 f32 D-regs. lane_div_N = lane//32, lane_mod_N = lane%32. -MFMA_M = 32 -MFMA_N = 32 -MFMA_K = 64 - - -def _mfma_scale_f32_32x32x64_fp8(result_type, operands): - """operands: [a, b, c] or [a, b, c, cbsz, blgp, opsel_a, scale_a, opsel_b, scale_b].""" - # Neutral e8m0 scale: 4 packed bytes of 127 => 2^(127-127) = 1.0 each. - neutral = arith.constant(0x7F7F7F7F, type=T.i32) - ops = list(operands) - if len(ops) < 9: - while len(ops) < 6: - ops.append(0) - ops.extend([neutral, 0, neutral]) - return rocdl.mfma_scale_f32_32x32x64_f8f6f4(result_type, ops) - - -MFMA_FN = _mfma_scale_f32_32x32x64_fp8 - -DEFAULT_COMPILE_HINTS = { - "waves_per_eu": 2, - "fast_fp_math": True, -} - - -# Via Uint32 because `//` on a signed Int32 lowers to floordivsi, which expands -# to several instructions for negative-operand rounding. Everything divided here -# is non-negative, so this keeps the single-instruction divui/remui. -def udiv(a, b): - return fx.Int32(fx.Uint32(a) // fx.Uint32(b)) - - -def umod(a, b): - return fx.Int32(fx.Uint32(a) % fx.Uint32(b)) - - -@lru_cache(maxsize=8) -def device_cu_count(device_index: int) -> int: - """Compute-unit count for a CUDA/HIP device (cached); 256 if unavailable.""" - try: - return torch.cuda.get_device_properties(device_index).multi_processor_count - except Exception: # noqa: BLE001 - return 256 - - -def load_pack_v8i32(i32_view, byte_off_i32, lane8): - """Load 32 fp8 bytes as the vector<8xi32> A/B operand of the 32x32x64 MFMA. - - Each lane holds 4 K-groups of 8 bytes at ``lane8 + kk*16``. Reads go out as - 2-dword loads because a ``v8i8`` buffer_load fails to lower. - """ - - def _load_i64(off): - v2 = i32_view.vec_load((udiv(byte_off_i32 + lane8 + off, 4),), vec_size=2) - return Vec(v2).bitcast(fx.Int64)[0].ir_value() - - i64_0 = _load_i64(0) # K-group 0: k = lane_div_N*8 + 0..7 - i64_1 = _load_i64(16) # K-group 1: k = lane_div_N*8 + 16..23 - i64_2 = _load_i64(32) # K-group 2: k = lane_div_N*8 + 32..39 - i64_3 = _load_i64(48) # K-group 3: k = lane_div_N*8 + 48..55 - return Vec.from_elements([i64_0, i64_1, i64_2, i64_3], fx.Int64).bitcast(fx.Int32) - - -def make_kv_key_view(kv_cache, head_size, kv_block_size, index_dim, preshuffle): - """Dword view of the co-packed KV keys over the gather's own axes: - ``(block, token//16, token%16, K-step, sub-group, lane half)``. - - The two data layouts differ only in this stride vector -- block-flat puts - token c at ``c*D``, shuffle_weight(16,16) at ``(c//16)*16D + (c%16)*16`` - with hidden groups every 256B -- so Preshuffle is a layout, not a second - gather path. index_dim is a multiple of 4, so the block stride divides. - """ - D = head_size - blk = udiv(fx.Int32(kv_block_size) * index_dim, 4) - stride = ( - (blk, D * 4, 4, 256, 64, 2) if preshuffle else (blk, D * 4, D // 4, 16, 4, 2) - ) - return GTensor(kv_cache, dtype=T.i32, shape=(-1,) * 6, stride=stride) - - -def make_kv_scale_view(kv_cache, head_size, kv_block_size, index_dim): - """f32 ``(block, token)`` view of the co-packed scales. The tail follows the - KVB*D key bytes and is token-ordered in both layouts, so it folds into base. - """ - return GTensor( - kv_cache, - dtype=T.f32, - shape=(-1, kv_block_size), - stride=(udiv(fx.Int32(kv_block_size) * index_dim, 4), 1), - base_offset=kv_block_size * head_size // 4, - ) - - -def load_pack_kv(key_view, physical, tok_in_block, kk_step, lane_div_N): - """vector<8xi32> B operand for one KV token and K-step. Layout-agnostic: - the view's strides carry the block-flat/preshuffle difference. - """ - c_hi, c_lo = udiv(tok_in_block, 16), umod(tok_in_block, 16) - # Slice everything but `sub` first: that fixes the token's base offset once - # instead of re-deriving it inside each of the four loads. - sub_view = key_view[physical, c_hi, c_lo, kk_step, None, lane_div_N] - - def _load_i64(sub): - v2 = sub_view.vec_load((sub,), vec_size=2) - return Vec(v2).bitcast(fx.Int64)[0].ir_value() - - return Vec.from_elements( - [_load_i64(sub) for sub in range_constexpr(4)], fx.Int64 - ).bitcast(fx.Int32) - - -def make_out_row_view(logits, stride_out, row_i32): - """1-D output GTensor for ``row_i32``, with the row byte base folded into - the pointer in i64 so the column offset stays i32. A 2-D (row, col) view - would compute ``row*stride + col`` in i32 and silently overflow past 2^31. - """ - stride_i64 = arith.extui(T.i64, _to_raw(stride_out)) - _ri64 = arith.extui(T.i64, _to_raw(row_i32)) - _byte = arith.muli(arith.muli(_ri64, stride_i64), arith.constant(4, type=T.i64)) - _idx = arith.index_cast(T.index, _byte) - return GTensor(logits, dtype=T.f32, shape=(-1,), static_bytes_offset_i64=_idx) - - -def mfma_head_reduce( - a_row, - b_col, - w_row, - kv_scale, - *, - m_tiles, - k_steps, - dreg_count=16, - mfma_fn=MFMA_FN, -): - """One column's logit: fp8 MFMA over heads, ReLU * weight sum, kv-scale, - in-wave head reduce. - - ``kv_scale`` (>= 0) is applied once to the column sum rather than per head: - ReLU is positively homogeneous, so ReLU(s*x) = s*ReLU(x) and hoisting it out - of the head sum is exact. - """ - res_ty = Vec.make_type(dreg_count, fx.Float32) - f32_0 = fx.Float32(0.0) - col_sum = f32_0 - for mi in range_constexpr(m_tiles): - acc = Vec.filled(dreg_count, 0.0, fx.Float32) - for kk in range_constexpr(k_steps): - acc = mfma_fn(res_ty, [a_row[mi][kk], b_col[kk], acc, 0, 0, 0]) - for ii in range_constexpr(dreg_count): - score = fx.Float32(Vec(acc)[ii]) - relu = score.maximumf(f32_0) - col_sum = col_sum + relu * w_row[mi][ii] - col_sum = col_sum * fx.Float32(kv_scale) - - # For 32x32 tile: lane_div_N = lane // 32 (0 or 1); shuffle_xor(32) sums them. - shuffles = (32,) if dreg_count == 16 else (16, 32) - for sh in shuffles: - peer = col_sum.shuffle_xor(sh, 64) - col_sum = col_sum + peer - return col_sum diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index e3e92aab497..374975f356a 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -1,674 +1,13 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Paged FP8 MQA logits (DeepSeek lightning indexer, decode) -- FlyDSL gfx950. +"""Paged FP8 MQA logits (decode) -- gfx950 16x16x128 mapping. -Compute for each decode batch element ``b`` and next-n query slot ``n`` (query -row ``b*next_n + n``) and KV logical position ``p``:: - - logits[b*next_n+n, p] = sum_h ReLU() * weights[b*next_n+n, h] * kv_scale_phys(p) - -masked by the causal rule ``p <= context_lens[b] - next_n + n`` (positions past -the query token and ``p >= context_lens[b]`` stay ``-inf``). ``K_phys(p)`` / -``kv_scale_phys(p)`` are gathered from a **paged** cache: logical position ``p`` -maps through the block table (``kv_indices``) to a physical block, and each -token's fp8 K bytes and its f32 dequant scale are **co-packed** in that block. - -Supports ``KVBlockSize >= 1`` with SplitKV KV-column parallelism, and mirrors -the tensor contract of the Triton ``deepgemm_fp8_paged_mqa_logits`` so the two -are interchangeable. ``Preshuffle=True`` consumes the production -``shuffle_weight(layout=(16,16))`` KV layout (needs ``KVBlockSize % 16 == 0``). - -The public ``flydsl_fp8_paged_mqa_logits`` entry dispatches the gfx950 -16x16x128 mapping (``fp8_paged_mqa_logits_gfx950``) for the production decode -shape ``H in {32,64}``, ``D=128``, ``KVBlockSize=64``, ``next_n in {1,2}``, -preshuffled cache. Other shapes keep this 32x32x64 generic kernel. Force the -generic path with ``FLYDSL_FP8_PAGED_MQA_LOGITS_GENERIC=1``. - -Cache layout, per physical block of ``KVBlockSize`` tokens: the fp8 key rows -(``D`` bytes each) grouped first, then the f32 dequant scales. Each lane -resolves its own column through ``kv_indices[b, p // KVBlockSize]``, so a column -tile may straddle blocks; ``KVBlockSize == 1`` degenerates to per-token -``[D fp8 | 4 scale]`` slots. - -The MFMA compute, ReLU*weight head-sum, head reduce, dword-pack load and output -view are shared with the dense kernel via ``._mqa_logits_common``. +The implementation lives in ``fp8_paged_mqa_logits_gfx950``. This module +re-exports the public entry so existing import paths keep working. """ -# No `from __future__ import annotations`: FlyDSL arg typing needs real -# annotation objects, not PEP 563 strings. - -import os -from functools import lru_cache - -import flydsl.compiler as flyc -import flydsl.expr as fx -import torch -from flydsl.expr import arith, range_constexpr -from flydsl.expr.typing import T - -from aiter.jit.utils.chip_info import get_gfx - -from ..tensor_shim import GTensor, _run_compiled, _to_raw -from ._mqa_logits_common import ( - DEFAULT_COMPILE_HINTS, - MFMA_K, - MFMA_M, - MFMA_N, - device_cu_count, - load_pack_kv, - load_pack_v8i32, - make_kv_key_view, - make_kv_scale_view, - make_out_row_view, - mfma_head_reduce, - udiv, - umod, +from .fp8_paged_mqa_logits_gfx950 import ( # noqa: F401 + flydsl_fp8_paged_mqa_logits, + flydsl_fp8_paged_mqa_logits_gfx950, ) - -_SIMDS_PER_CU = 4 - - -def _auto_split_kv( - batch_size, next_n, wave_per_eu, device_index, total_cu=None, waves_per_block=1 -): - """Pick SplitKV so the dispatch lands ``wave_per_eu`` waves on every SIMD: - - SplitKV = round(WavePerEU * 4 * TotalCuCount / (batch*next_n*WavePerBlock)) - - Resident wave count is the lever that matters here: the kernel is bound by - memory throughput rather than per-wave latency, so too few waves leaves the - memory system idle and too many just queue. Accounting for the CTA's wave - count keeps the target fixed across shapes. Returns >= 1. - """ - tile_q_count = max(1, batch_size * next_n) - if total_cu is None: - total_cu = device_cu_count(device_index) - target_waves = max(1, wave_per_eu) * _SIMDS_PER_CU * total_cu - split_kv = round(target_waves / (tile_q_count * max(1, waves_per_block))) - return max(1, int(split_kv)) - - -# Default KV tile width (columns processed per MFMA inner-loop iteration). -_BLOCK_KV = 128 - - -def _uceildiv(a, b): - a = fx.Int32(a) - b = fx.Int32(b) - return fx.Int32((fx.Uint32(a) + fx.Uint32(b) - 1) // fx.Uint32(b)) - - -def _imin(a, b): - a = fx.Int32(a) - b = fx.Int32(b) - return (a <= b).select(a, b) - - -def _build_paged_kernel( - *, - num_heads: int, - head_size: int, - block_kv: int, - waves_per_block: int, - kv_block_size: int = 1, - preshuffle: bool = False, -): - """Paged MQA-logits kernel with SplitKV column parallelism. - - One thread block owns one ``(batch, next_n)`` query row and walks its KV - window in ``BKV``-wide tiles; the ``waves_per_block`` waves each take a - disjoint slice of the ``BKV/MFMA_N`` column tiles, with no barrier between - them. - - ``preshuffle`` only changes the intra-block key-byte offset -- the - block-table gather, scale gather, MFMA and mask are identical either way. - """ - H = num_heads - D = head_size - BKV = block_kv - WPB = waves_per_block - KVB = kv_block_size - assert KVB >= 1, f"kv_block_size must be >= 1; got {KVB}" - MR_BLOCK_THREADS = 64 * WPB - - assert H % MFMA_M == 0, f"num_heads={H} must be a multiple of MFMA_M={MFMA_M}" - assert BKV % MFMA_N == 0, f"block_kv={BKV} must be a multiple of MFMA_N={MFMA_N}" - assert D % MFMA_K == 0, f"head_size={D} must be a multiple of MFMA_K={MFMA_K}" - assert WPB >= 1, "waves_per_block must be >= 1" - N_TILES = BKV // MFMA_N - assert ( - N_TILES % WPB == 0 - ), f"BKV/MFMA_N={N_TILES} must be divisible by waves_per_block={WPB}" - M_TILES = H // MFMA_M - K_STEPS = D // MFMA_K - N_TILES_PER_WAVE = N_TILES // WPB - - _kname = ( - f"fp8_paged_mqa_logits_H{H}_D{D}_bkv{BKV}_kvb{KVB}_w{WPB}" - f"{'_ps' if preshuffle else ''}_flydsl" - ) - - @flyc.kernel(name=_kname, known_block_size=[MR_BLOCK_THREADS, 1, 1]) - def kernel( - Q: fx.Tensor, # [batch, next_n, H, D] fp8 (bytes passed raw) - KV_cache: fx.Tensor, # [num_blocks, KVB, 1, index_dim] uint8 block-flat - weights: fx.Tensor, # [batch*next_n, H] f32 - out_logits: fx.Tensor, # [batch*next_n, max_model_len] f32 (-inf prefilled) - context_lens: fx.Tensor, # [batch] i32 - kv_indices: fx.Tensor, # [batch, max_block_len] i32 (block table) - next_n: fx.Int32, - batch_size: fx.Int32, - split_kv: fx.Int32, # KV-column splits (grid = split_kv*batch*next_n) - stride_q_batch: fx.Int32, # fp8 elems (== bytes) - stride_q_next_n: fx.Int32, - stride_q_heads: fx.Int32, - index_dim: fx.Int32, # per-token slot bytes (D+4+pad); block stride = KVB*index_dim - max_block_len: fx.Int32, # kv_indices row width - stride_out: fx.Int32, # out_logits.stride(0) == max_model_len - ): - tid = fx.thread_idx.x - bid = fx.block_idx.x - - # ChunkQ==H => grid = split_kv * batch * next_n (split outermost), - # mirroring the Triton _deepgemm_fp8_paged_mqa_logits pid decomposition. - pid_next_n = umod(bid, next_n) - _rem = udiv(bid, next_n) - pid_batch = umod(_rem, batch_size) - pid_split_kv = udiv(_rem, batch_size) - - wave = udiv(tid, 64) - lane = umod(tid, 64) - lane_div_N = udiv(lane, MFMA_N) - lane_mod_N = umod(lane, MFMA_N) - lane8 = lane_div_N * 8 - - q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) - # Uniform-base views: the per-token gather rides the per-lane byte - # offset, not the buffer base, since a per-lane base cannot ride a - # scalar buffer descriptor. That offset stays i32, hence the host's - # pool-size assert. - kv_key = make_kv_key_view(KV_cache, D, KVB, index_dim, preshuffle) - kv_scale_v = make_kv_scale_view(KV_cache, D, KVB, index_dim) - w_t = GTensor(weights, dtype=T.f32, shape=(-1, H)) - cl_t = GTensor(context_lens, dtype=T.i32, shape=(-1,)) - ind_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) - - context_length = fx.Int32(cl_t[pid_batch]) - # Inclusive causal upper bound: p <= context_length - next_n + pid_next_n. - q_limit = context_length - next_n + pid_next_n - - out_row = pid_batch * next_n + pid_next_n - out_row_t = make_out_row_view(out_logits, stride_out, out_row) - - # ---- Preload Q frags + weights for the single query row ---- - # The A-operand layout is per in-wave lane, so `lane` indexes Q, not - # `tid`. - q_row_base = pid_batch * stride_q_batch + pid_next_n * stride_q_next_n - a_pack = [[None] * K_STEPS for _ in range_constexpr(M_TILES)] - for mi in range_constexpr(M_TILES): - h_a = mi * MFMA_M + lane_mod_N - base_a = q_row_base + h_a * stride_q_heads - for kk in range_constexpr(K_STEPS): - byte_base = base_a + kk * MFMA_K - a_pack[mi][kk] = load_pack_v8i32(q_i32, byte_base, lane8) - - # weights[out_row, h] per (mi, ii): head = mi*MFMA_M + h_off(lane_div_N, ii) - # For the 32x32x64 MFMA (16 f32 D-regs), the D-reg ii at lane_div_N encodes - # head: h_off = (ii % 4) + lane_div_N * 4 + (ii // 4) * 8 - # -- the same 4-interleave as the 16x16x32 layout, extended to 16 D-regs. - _DREG = 16 - w_frag = [[None] * _DREG for _ in range_constexpr(M_TILES)] - for mi in range_constexpr(M_TILES): - for ii in range_constexpr(_DREG): - h_off = (ii % 4) + lane_div_N * 4 + (ii // 4) * 8 - h_w = mi * MFMA_M + h_off - w_frag[mi][ii] = fx.Float32(w_t[out_row, h_w]) - - # ---- SplitKV: each split owns ceil(total_tiles/split_kv) contiguous - # BKV-aligned tiles, so the slices are gap-free and disjoint and every - # logits column has exactly one writer -- no cross-CTA reduction. A - # split past ctx runs zero iterations. ---- - context_chunk_num = _uceildiv(context_length, fx.Int32(BKV)) - split_chunk_num = _uceildiv(context_chunk_num, split_kv) - full_end = context_chunk_num * BKV - split_cols = split_chunk_num * BKV - tile_lo_col = pid_split_kv * split_cols - tile_hi_col = _imin(tile_lo_col + split_cols, full_end) - ctx_m1 = context_length - 1 - - for col0 in range(tile_lo_col, tile_hi_col, fx.Int32(BKV)): - wave_ni_base = wave * N_TILES_PER_WAVE - - # Helpers close over this chunk's SSA values; default args silence B023. - def _col_for(abs_ni, _col0=col0): - return _col0 + abs_ni * MFMA_N + lane_mod_N - - def _col_c_for(abs_ni, _ctx_m1=ctx_m1): - return _imin(_col_for(abs_ni), _ctx_m1) - - def _ind_off_for(abs_ni): - return pid_batch * max_block_len + udiv(_col_c_for(abs_ni), KVB) - - def _tok_in_block_for(abs_ni): - return umod(_col_c_for(abs_ni), KVB) - - def _load_physical(ind_off): - return fx.Int32(ind_t[ind_off]) - - def _prefetch_physical(physical_cur, ind_off_cur, ind_off_next): - if ind_off_next == ind_off_cur: - return physical_cur - return _load_physical(ind_off_next) - - def _b_col_for_tile(physical, tok_in_block, b0_carried): - return [ - b0_carried, - load_pack_kv(kv_key, physical, tok_in_block, 1, lane_div_N), - ] - - def _process_tile_with_b0(abs_ni, physical, b0_carried): - col = _col_for(abs_ni) - tok_in_block = _tok_in_block_for(abs_ni) - b_col = _b_col_for_tile(physical, tok_in_block, b0_carried) - kv_scale = fx.Float32(kv_scale_v[physical, tok_in_block]) - col_sum = mfma_head_reduce( - a_pack, b_col, w_frag, kv_scale, m_tiles=M_TILES, k_steps=K_STEPS - ) - is_writer = (lane_div_N == 0) & (col <= q_limit) - - def _do_write(_t=out_row_t, _c=col, _v=col_sum): - _t[_c] = _v - - @flyc.jit - def _guarded_write(_pred=is_writer, _w=_do_write): - if _pred: - _w() - - _guarded_write() - - def _prefetch_k0(abs_ni, physical): - return load_pack_kv( - kv_key, physical, _tok_in_block_for(abs_ni), 0, lane_div_N - ) - - physical_carry = _load_physical(_ind_off_for(wave_ni_base + 0)) - b0_carry = _prefetch_k0(wave_ni_base + 0, physical_carry) - for ni in range_constexpr(N_TILES_PER_WAVE): - abs_ni = wave_ni_base + ni - physical = physical_carry - if ni + 1 < N_TILES_PER_WAVE: - abs_ni_next = wave_ni_base + ni + 1 - physical_carry = _prefetch_physical( - physical, - _ind_off_for(abs_ni), - _ind_off_for(abs_ni_next), - ) - _process_tile_with_b0(abs_ni, physical, b0_carry) - if ni + 1 < N_TILES_PER_WAVE: - abs_ni_next = wave_ni_base + ni + 1 - b0_carry = _prefetch_k0(abs_ni_next, physical_carry) - - @flyc.jit - def launch_fp8_paged_mqa_logits( - Q: fx.Tensor, - KV_cache: fx.Tensor, - weights: fx.Tensor, - out_logits: fx.Tensor, - context_lens: fx.Tensor, - kv_indices: fx.Tensor, - grid_blocks: fx.Int32, - next_n: fx.Int32, - batch_size: fx.Int32, - split_kv: fx.Int32, - stride_q_batch: fx.Int32, - stride_q_next_n: fx.Int32, - stride_q_heads: fx.Int32, - index_dim: fx.Int32, - max_block_len: fx.Int32, - stride_out: fx.Int32, - stream: fx.Stream = fx.Stream(None), # noqa: B008 - ): - gx = arith.index_cast(T.index, _to_raw(grid_blocks)) - kernel._func.__name__ = _kname - kernel( - Q, - KV_cache, - weights, - out_logits, - context_lens, - kv_indices, - next_n, - batch_size, - split_kv, - stride_q_batch, - stride_q_next_n, - stride_q_heads, - index_dim, - max_block_len, - stride_out, - ).launch(grid=(gx, 1, 1), block=(MR_BLOCK_THREADS, 1, 1), stream=stream) - - return launch_fp8_paged_mqa_logits - - -# Kernel variants: single-token-per-block, wave-split only ("paged_w"). -# WPB must divide the column-tile count BKV/16 (=8 at the default BKV=128). -# -# WPB=1 (one wave per CTA) is the default: a 1-wave CTA schedules more freely -# and pays no whole-CTA granularity cost, and the waves have no work to share -# anyway, each owning a disjoint set of KV columns with no barrier or LDS -# between them. -KERNEL_VARIANTS = tuple(f"paged_w{w}" for w in (1, 2, 4)) -DEFAULT_VARIANT = "paged_w1" - - -def _variant_wpb(variant): - """Parse the WPB int out of a ``paged_w`` tag (single validation point).""" - if variant not in KERNEL_VARIANTS: - raise ValueError( - f"unknown fp8_paged_mqa_logits variant {variant!r}; " - f"available: {list(KERNEL_VARIANTS)}" - ) - return int(variant.removeprefix("paged_w")) - - -def _resolve_variant(variant): - """Pick the variant tag: explicit arg, then env override, then default. - Shape-adaptive selection is not implemented.""" - return ( - variant - or os.environ.get("FLYDSL_FP8_PAGED_MQA_LOGITS_VARIANT") - or DEFAULT_VARIANT - ) - - -@lru_cache(maxsize=32) -def compile_fp8_paged_mqa_logits( - *, - num_heads: int, - head_size: int, - block_kv: int = _BLOCK_KV, - kv_block_size: int = 1, - preshuffle: bool = False, - variant: str = DEFAULT_VARIANT, -): - """Return a cached, compiled FlyDSL paged launcher for the given config. - - ``num_heads``/``head_size``/``kv_block_size``/``preshuffle`` are compile-time - constants and ``variant`` is a ``paged_w`` tag. - """ - launcher = _build_paged_kernel( - num_heads=num_heads, - head_size=head_size, - block_kv=block_kv, - waves_per_block=_variant_wpb(variant), - kv_block_size=kv_block_size, - preshuffle=preshuffle, - ) - launcher.compile_hints = dict(DEFAULT_COMPILE_HINTS) - return launcher - - -def _use_gfx950_fastpath( - q_fp8, - kv_cache, - *, - Preshuffle, - KVBlockSize, - ChunkK, - WavePerEU, - variant, -): - """True for the production decode shape the 16x16x128 gfx950 kernel covers.""" - if os.environ.get("FLYDSL_FP8_PAGED_MQA_LOGITS_GENERIC", "").strip() in ( - "1", - "true", - "True", - ): - return False - env_variant = os.environ.get("FLYDSL_FP8_PAGED_MQA_LOGITS_VARIANT") - if variant not in (None, DEFAULT_VARIANT) or ( - env_variant and env_variant != DEFAULT_VARIANT - ): - return False - if not Preshuffle or int(KVBlockSize) != 64: - return False - if int(ChunkK) != _BLOCK_KV or int(WavePerEU) != 2: - return False - if get_gfx() != "gfx950": - return False - _, next_n, num_heads, head_size = q_fp8.shape - if num_heads not in (32, 64) or head_size != 128 or next_n not in (1, 2): - return False - from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype - - if q_fp8.dtype != get_fp8_e4m3_dtype(): - return False - _, block_size, one, index_dim = kv_cache.shape - return block_size == 64 and one == 1 and index_dim == 132 - - -def flydsl_fp8_paged_mqa_logits_generic( - q_fp8, - kv_cache, - weights, - out_logits, - context_lens, - kv_indices, - max_model_len, - *, - Preshuffle=False, - KVBlockSize=1, - ChunkK=_BLOCK_KV, - SplitKV=None, - WavePerEU=2, - TotalCuCount=None, - variant=None, - stream=None, -): - """Generic 32x32x64 FlyDSL paged FP8 MQA logits (decode). - - Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. - Production H32/H64 D128 KVB64 preshuffle decode uses the gfx950 mapping via - ``flydsl_fp8_paged_mqa_logits`` instead. - - q_fp8: [batch, next_n, heads, hidden_dim], float8 e4m3fn (gfx950 native) - kv_cache: [num_blocks, KVBlockSize, 1, index_dim] uint8, co-packed per - block as KVBlockSize*hidden_dim fp8 K bytes then KVBlockSize - f32 scales; index_dim == hidden_dim + 4 (+ optional padding). - Preshuffle: consume the shuffle_weight(16x16) layout for the fp8 data - section (the scale tail stays token-ordered). Requires - ``KVBlockSize % 16 == 0``. - weights: [batch*next_n, heads], float32 - out_logits: [batch*next_n, max_model_len], float32. MUST be prefilled with - -inf: the kernel writes only in-window positions. - context_lens: [batch], int32 - kv_indices: [batch, max_block_len], int32 block table indexed by - ``p // KVBlockSize``. - max_model_len: out_logits column count. - SplitKV: KV-column split count (grid = split_kv*batch*next_n). None => - sized for WavePerEU waves per SIMD; 1 disables splitting. - WavePerEU: target resident waves per SIMD for the auto-SplitKV formula. - TotalCuCount: CU count for that formula; None => query the device. - ChunkK: KV tile width, a multiple of MFMA_N=32. Defaults to 128. - - Returns the same ``out_logits`` tensor (written in place). - """ - if Preshuffle: - assert KVBlockSize % 16 == 0, ( - f"Preshuffle mode only supports KVBlockSize aligned to 16; " - f"got KVBlockSize={KVBlockSize}." - ) - if ChunkK % MFMA_N != 0: - raise ValueError(f"ChunkK={ChunkK} must be a multiple of MFMA_N={MFMA_N}.") - - batch_size, next_n, num_heads, head_size = q_fp8.shape - assert num_heads & (num_heads - 1) == 0, "num q. heads should be power of 2." - assert head_size & (head_size - 1) == 0, "head size should be power of 2." - - num_blocks, block_size, one, index_dim = kv_cache.shape - assert block_size == KVBlockSize, ( - f"kv_cache KVBlockSize dim ({block_size}) must equal KVBlockSize " - f"({KVBlockSize})." - ) - assert KVBlockSize >= 1, f"KVBlockSize must be >= 1; got {KVBlockSize}." - assert one == 1, f"kv_cache head dim must be 1; got {one}." - assert ( - index_dim >= head_size + 4 - ), f"index_dim={index_dim} must hold {head_size} fp8 bytes + 4 scale bytes." - assert ( - kv_cache.dtype == torch.uint8 - ), f"kv_cache must be uint8 co-packed bytes; got {kv_cache.dtype}." - - # i32 gather-offset ceiling: the per-token byte base (physical block byte base - # + intra-block offset) is computed in i32 (it rides the buffer voffset), so - # the whole cache pool must be addressable in i32 bytes. A larger pool would - # need an i64/global gather. - pool_bytes = num_blocks * block_size * index_dim - assert pool_bytes < 2**31, ( - f"num_blocks*KVBlockSize*index_dim={pool_bytes} exceeds the i32 " - f"gather-offset limit (2^31); the paged kernel needs an i64 gather path " - f"for a cache pool this large." - ) - - _, max_block_len = kv_indices.shape - - # FlyDSL's DLPack adaptor rejects 0-dim tensors; keep logical ranks. - context_lens = context_lens.reshape(batch_size) - weights = weights.reshape(batch_size * next_n, num_heads) - kv_indices = kv_indices.reshape(batch_size, max_block_len) - - _fnuz = torch.float8_e4m3fnuz - _fn = torch.float8_e4m3fn - assert q_fp8.dtype in ( - _fnuz, - _fn, - ), f"q_fp8 must be e4m3 fp8 (fnuz or fn); got {q_fp8.dtype}" - assert ( - get_gfx() == "gfx950" - ), f"flydsl_fp8_paged_mqa_logits_generic targets gfx950 (32x32x64 MFMA); got {get_gfx()}" - - variant = _resolve_variant(variant) - - launcher = compile_fp8_paged_mqa_logits( - num_heads=num_heads, - head_size=head_size, - block_kv=ChunkK, - kv_block_size=int(KVBlockSize), - preshuffle=bool(Preshuffle), - variant=variant, - ) - - # Co-packed cache -> raw uint8 byte view [num_blocks * KVBlockSize * index_dim]. - kv_bytes = kv_cache.reshape(-1) - - # KV-column splits (one lever): fill the device when the batch*next_n row - # grid is small. Auto targets WavePerEU waves per SIMD; overridable. - if SplitKV is None: - split_kv = _auto_split_kv( - batch_size, - next_n, - WavePerEU, - q_fp8.device.index, - TotalCuCount, - waves_per_block=_variant_wpb(variant), - ) - else: - split_kv = max(1, int(SplitKV)) - - grid_blocks = batch_size * next_n * split_kv - - if stream is None: - stream = torch.cuda.current_stream() - - with torch.cuda.device(q_fp8.device.index): - _run_compiled( - launcher, - q_fp8, - kv_bytes, - weights, - out_logits, - context_lens, - kv_indices, - int(grid_blocks), - int(next_n), - int(batch_size), - int(split_kv), - int(q_fp8.stride(0)), - int(q_fp8.stride(1)), - int(q_fp8.stride(2)), - int(index_dim), - int(max_block_len), - int(out_logits.stride(0)), - stream, - ) - - return out_logits - - -def flydsl_fp8_paged_mqa_logits( - q_fp8, - kv_cache, - weights, - out_logits, - context_lens, - kv_indices, - max_model_len, - *, - Preshuffle=False, - KVBlockSize=1, - ChunkK=_BLOCK_KV, - SplitKV=None, - WavePerEU=2, - TotalCuCount=None, - variant=None, - stream=None, -): - """FlyDSL paged FP8 MQA logits (decode) -- KVBlockSize>=1 with SplitKV. - - Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract. - Dispatches the gfx950 16x16x128 kernel for H32/H64 D128 KVB64 preshuffle - ``next_n in {1,2}``; otherwise the generic 32x32x64 kernel. - """ - if _use_gfx950_fastpath( - q_fp8, - kv_cache, - Preshuffle=Preshuffle, - KVBlockSize=KVBlockSize, - ChunkK=ChunkK, - WavePerEU=WavePerEU, - variant=variant, - ): - from .fp8_paged_mqa_logits_gfx950 import flydsl_fp8_paged_mqa_logits_gfx950 - - return flydsl_fp8_paged_mqa_logits_gfx950( - q_fp8, - kv_cache, - weights, - out_logits, - context_lens, - kv_indices, - max_model_len, - KVBlockSize=KVBlockSize, - SplitKV=SplitKV, - TotalCuCount=TotalCuCount, - stream=stream, - ) - return flydsl_fp8_paged_mqa_logits_generic( - q_fp8, - kv_cache, - weights, - out_logits, - context_lens, - kv_indices, - max_model_len, - Preshuffle=Preshuffle, - KVBlockSize=KVBlockSize, - ChunkK=ChunkK, - SplitKV=SplitKV, - WavePerEU=WavePerEU, - TotalCuCount=TotalCuCount, - variant=variant, - stream=stream, - ) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 30159d0fe4b..267269ac594 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -3,8 +3,8 @@ """gfx950 H32/H64 D128/KVB64 preshuffled paged FP8 MQA-logits kernel. -Default decode path for ``flydsl_fp8_paged_mqa_logits`` on this shape. -``flydsl_fp8_paged_mqa_logits_gfx950`` remains the explicit entry. +Public decode path for ``flydsl_fp8_paged_mqa_logits``. Supports H in {32, 64}, +``next_n`` in {1, 2}, and KVBlockSize=64 with a preshuffled KV cache. """ from functools import lru_cache @@ -12,7 +12,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl.expr import arith, range_constexpr +from flydsl.expr import arith, range_constexpr, rocdl from flydsl.expr.typing import T from aiter.jit.utils.chip_info import get_gfx @@ -20,45 +20,181 @@ from .. import buffer_ops from ..tensor_shim import GTensor, _run_compiled, _to_raw -from ._fp8_paged_mqa_logits_gfx950 import ( - B_RING, - DREG, - HEAD_DIM, - KV_BLOCK_SIZE, - MFMA_M, - MFMA_N, - PAGE_VMEM_LOADS, - SUPPORTED_HEADS, - guarded_store, - imin, - load_kv_scale, - load_preshuffled_k_pack, - load_q_pack, - mfma_scores, - page_vmem_stores, - reduce_scores, - schedule_mfma_valu_pairs, - uceildiv, - wait_vmcnt, -) -from ._mqa_logits_common import ( - DEFAULT_COMPILE_HINTS, - device_cu_count, - udiv, - umod, -) + +Vec = fx.Vector + +SUPPORTED_HEADS = (32, 64) +HEAD_DIM = 128 +INDEX_DIM = HEAD_DIM + 4 +KV_BLOCK_SIZE = 64 +MFMA_M = MFMA_N = 16 +DREG = 4 +B_RING = KV_BLOCK_SIZE // MFMA_N +BLOCK_I32 = KV_BLOCK_SIZE * INDEX_DIM // 4 +SCALE_I32 = KV_BLOCK_SIZE * HEAD_DIM // 4 +TILE_I32 = HEAD_DIM * MFMA_N // 4 +HALF_I32 = MFMA_N * 16 // 4 +LANE_K_I32 = (HEAD_DIM // (64 // MFMA_N)) * MFMA_N // 4 +# Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). +PAGE_VMEM_LOADS = B_RING * 3 +_NEUTRAL_E8M0 = 0x7F7F7F7F _GFX950 = "gfx950" +DEFAULT_COMPILE_HINTS = { + "waves_per_eu": 2, + "fast_fp_math": True, +} + + +def wait_vmcnt(n): + """Explicit s_waitcnt vmcnt(n). Leaves n VM ops outstanding.""" + rocdl.sched_barrier(0) + rocdl.s_waitcnt(vmcnt=int(n)) + rocdl.sched_barrier(0) + + +# Via Uint32 because `//` on a signed Int32 lowers to floordivsi. +def udiv(a, b): + return fx.Int32(fx.Uint32(a) // fx.Uint32(b)) + + +def umod(a, b): + return fx.Int32(fx.Uint32(a) % fx.Uint32(b)) -def _unwrap(value): - return value.ir_value() if hasattr(value, "ir_value") else value +def uceildiv(a, b): + a, b = fx.Int32(a), fx.Int32(b) + return fx.Int32((fx.Uint32(a) + fx.Uint32(b) - 1) // fx.Uint32(b)) + + +def imin(a, b): + a, b = fx.Int32(a), fx.Int32(b) + return (a <= b).select(a, b) + + +def guarded_store(pred, store_fn): + """Predicated side effect; masked buffer stores can fault on gfx950.""" + + @flyc.jit + def _guarded(_pred=pred, _store=store_fn): + if _pred: + _store() + + _guarded() -def _build_kernel(*, index_dim: int, num_heads: int, next_n: int): +def _concat_i32x4(lo, hi): + return Vec(lo).shuffle(Vec(hi), list(range(DREG * 2))) + + +def load_q_pack(q_i32, byte_base, lane_div_16): + """Load one row's lane-owned contiguous K32 segment as a 256-bit operand.""" + off = (byte_base + lane_div_16 * 32) // fx.Int32(4) + lo = buffer_ops.buffer_load(q_i32.rsrc, off, vec_width=4, dtype=T.i32) + hi = buffer_ops.buffer_load(q_i32.rsrc, off + 4, vec_width=4, dtype=T.i32) + return _concat_i32x4(lo, hi) + + +def load_preshuffled_k_pack( + kv_i32, physical, tile_in_page, lane_mod_16, lane_div_16 +): + """Load a shuffle_weight(16,16) K column as an i32x8 MFMA operand.""" + base = ( + physical * BLOCK_I32 + + fx.Int32(tile_in_page * TILE_I32) + + lane_div_16 * LANE_K_I32 + + lane_mod_16 * 4 + ) + lo = buffer_ops.buffer_load( + kv_i32.rsrc, base, vec_width=4, dtype=T.i32, cache_modifier=2 + ) + hi = buffer_ops.buffer_load( + kv_i32.rsrc, + base + HALF_I32, + vec_width=4, + dtype=T.i32, + cache_modifier=2, + ) + return _concat_i32x4(lo, hi) + + +def load_kv_scale(kv_i32, physical, token_in_page): + off = physical * BLOCK_I32 + SCALE_I32 + token_in_page + return fx.Float32( + buffer_ops.buffer_load( + kv_i32.rsrc, off, vec_width=1, dtype=T.f32, cache_modifier=2 + ) + ) + + +def mfma_scores(mma, a_tiles, b_pack, *, m_tiles): + """Return one 4-f32 score fragment per 16-head M tile.""" + neutral = fx.Int32(_NEUTRAL_E8M0) + b_frag = fx.make_rmem_tensor(DREG * 2, fx.Int32) + b_frag.store(b_pack) + scores = [] + for mi in range_constexpr(m_tiles): + a_frag = fx.make_rmem_tensor(DREG * 2, fx.Int32) + acc = fx.make_rmem_tensor(DREG, fx.Float32) + a_frag.store(a_tiles[mi]) + acc.store(Vec.filled(DREG, 0.0, fx.Float32)) + fx.gemm( + mma, + acc, + a_frag, + b_frag, + acc, + scale_a=neutral, + scale_b=neutral, + ) + scores.append(acc.load()) + return scores + + +def reduce_scores(scores, weights, kv_scale, *, m_tiles): + """ReLU, weighted H reduction, positive KV scale, then wave reduction.""" + zero = fx.Float32(0.0) + total = zero + for mi in range_constexpr(m_tiles): + frag = Vec(scores[mi]) + for ii in range_constexpr(DREG): + total = total + fx.Float32(frag[ii]).maximumf(zero) * weights[mi][ii] + total = total * kv_scale + total = total + total.shuffle_xor(16, 64) + return total + total.shuffle_xor(32, 64) + + +def schedule_mfma_valu_pairs(*, m_tiles): + """Pair each 32-cycle MFMA with the prior tile's 12-op VALU fragment.""" + for _ in range_constexpr(m_tiles): + rocdl.sched_mfma(1) + rocdl.sched_group_barrier("valu", DREG * 3, 0) + + +@lru_cache(maxsize=8) +def device_cu_count(device_index: int) -> int: + try: + return torch.cuda.get_device_properties(device_index).multi_processor_count + except Exception: # noqa: BLE001 + return 256 + + +def _make_out_row_view(logits, stride_out, row): + """Fold the row byte offset into the pointer using i64 arithmetic.""" + byte = fx.Int64(fx.Uint32(row)) * fx.Int64(fx.Uint32(stride_out)) * 4 + return GTensor( + logits, + dtype=T.f32, + shape=(-1,), + cache_modifier=2, + static_bytes_offset_i64=byte, + ) + + +def _build_kernel(*, num_heads: int, next_n: int): m_tiles = num_heads // MFMA_M n_rows = int(next_n) - store_vmem = page_vmem_stores(n_rows) + store_vmem = n_rows kernel_name = ( f"fp8_paged_mqa_logits_gfx950_H{num_heads}_D{HEAD_DIM}_" f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{n_rows}_db2_ws1_pfp_sb0_nt_ps_flydsl" @@ -86,13 +222,17 @@ def kernel( lane = umod(fx.thread_idx.x, 64) lane_mod_16 = umod(lane, MFMA_N) lane_div_16 = udiv(lane, MFMA_N) + mma = fx.make_mma_atom( + fx.rocdl.cdna4.MFMA_Scale( + MFMA_M, MFMA_N, HEAD_DIM, fx.Float8E4M3FN + ) + ) q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) kv_i32 = GTensor(KV_cache, dtype=T.i32, shape=(-1,)) weight_t = GTensor(weights, dtype=T.f32, shape=(-1,)) context_t = GTensor(context_lens, dtype=T.i32, shape=(-1,)) table_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) - out_t = GTensor(out_logits, dtype=T.f32, shape=(-1,), cache_modifier=2) context_len = fx.Int32(context_t[pid_batch]) page_count = uceildiv(context_len, fx.Int32(KV_BLOCK_SIZE)) @@ -103,11 +243,17 @@ def kernel( col_hi = page_hi * KV_BLOCK_SIZE next_n_c = fx.Int32(n_rows) + out_rows = [ + _make_out_row_view( + out_logits, stride_out, pid_batch * next_n_c + fx.Int32(row) + ) + for row in range_constexpr(n_rows) + ] a_rows = [ [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(n_rows) ] w_rows = [ - [[None for _ in range_constexpr(DREG)] for _ in range_constexpr(m_tiles)] + [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(n_rows) ] for row in range_constexpr(n_rows): @@ -118,7 +264,7 @@ def kernel( h = mi * MFMA_M + lane_mod_16 byte_base = q_base + h * stride_q_heads a_rows[row][mi] = load_q_pack(q_i32, byte_base, lane_div_16) - weight_vec = fx.Vector( + w_rows[row][mi] = fx.Vector( buffer_ops.buffer_load( weight_t.rsrc, out_row * num_heads + mi * MFMA_M + lane_div_16 * DREG, @@ -126,8 +272,6 @@ def kernel( dtype=T.f32, ) ) - for ii in range_constexpr(DREG): - w_rows[row][mi][ii] = fx.Float32(weight_vec[ii]) def _load_physical(page_col): table_idx = pid_batch * max_block_len + udiv( @@ -150,17 +294,11 @@ def _issue_physical(physical): slot, lane_mod_16, lane_div_16, - index_dim=index_dim, ) ) - scale_slots.append( - load_kv_scale(kv_i32, physical, token, index_dim=index_dim) - ) + scale_slots.append(load_kv_scale(kv_i32, physical, token)) return b_slots, scale_slots - def _issue_page(page_col): - return _issue_physical(_load_physical(page_col)) - def _make_page_bank(): return ( [ @@ -175,7 +313,9 @@ def _store_page_bank(bank, b_slots, scale_slots): for slot in range_constexpr(B_RING): b_bank[slot].store(b_slots[slot]) scale_bank[slot].store( - fx.Vector.from_elements([_unwrap(scale_slots[slot])], fx.Float32) + fx.Vector.from_elements( + [scale_slots[slot].ir_value()], fx.Float32 + ) ) def _load_page_bank(bank): @@ -192,26 +332,22 @@ def _make_physical_bank(): return fx.make_rmem_tensor(1, fx.Int32) def _store_physical(bank, physical): - bank.store(fx.Vector.from_elements([_unwrap(physical)], fx.Int32)) + bank.store(fx.Vector.from_elements([physical.ir_value()], fx.Int32)) def _load_physical_bank(bank): return fx.Int32(fx.Vector(bank.load())[0]) - def _reduce_row(row, scores, scale): - return reduce_scores(scores, w_rows[row], scale, m_tiles=m_tiles) - def _store_page_row(row, page_col, logits): q_row = fx.Int32(row) q_limit = context_len - next_n_c + q_row - out_row = pid_batch * next_n_c + q_row col = page_col + lane logit = logits[0] for slot in range_constexpr(1, B_RING): logit = (lane_div_16 == slot).select(logits[slot], logit) writer = (col < context_len) & (col <= q_limit) - def _write(_row=out_row, _col=col, _value=logit): - out_t[_row * stride_out + _col] = _value + def _write(_out=out_rows[row], _col=col, _value=logit): + _out[_col] = _value guarded_store(writer, _write) @@ -223,28 +359,36 @@ def _compute_page(page_col, b_slots, scale_slots): ] for slot in range_constexpr(B_RING): scores = [ - mfma_scores(a_rows[row], b_slots[slot], m_tiles=m_tiles) + mfma_scores( + mma, a_rows[row], b_slots[slot], m_tiles=m_tiles + ) for row in range_constexpr(n_rows) ] if slot > 0: schedule_mfma_valu_pairs(m_tiles=m_tiles) for row in range_constexpr(n_rows): - page_logits[row][slot - 1] = _reduce_row( - row, prev_scores[row], prev_scale + page_logits[row][slot - 1] = reduce_scores( + prev_scores[row], + w_rows[row], + prev_scale, + m_tiles=m_tiles, ) prev_scores = scores prev_scale = scale_slots[slot] schedule_mfma_valu_pairs(m_tiles=m_tiles) for row in range_constexpr(n_rows): - page_logits[row][B_RING - 1] = _reduce_row( - row, prev_scores[row], prev_scale + page_logits[row][B_RING - 1] = reduce_scores( + prev_scores[row], + w_rows[row], + prev_scale, + m_tiles=m_tiles, ) _store_page_row(row, page_col, page_logits[row]) if page_lo < page_hi: current_bank = _make_page_bank() next_bank = _make_page_bank() - b_init, scale_init = _issue_page(col_lo) + b_init, scale_init = _issue_physical(_load_physical(col_lo)) _store_page_bank(current_bank, b_init, scale_init) current_physical = _make_physical_bank() next_physical = _make_physical_bank() @@ -337,11 +481,11 @@ def launch( @lru_cache(maxsize=8) -def _compile(*, index_dim: int, num_heads: int, next_n: int): - return _build_kernel(index_dim=index_dim, num_heads=num_heads, next_n=next_n) +def _compile(*, num_heads: int, next_n: int): + return _build_kernel(num_heads=num_heads, next_n=next_n) -def flydsl_fp8_paged_mqa_logits_gfx950( +def flydsl_fp8_paged_mqa_logits( q_fp8, kv_cache, weights, @@ -350,14 +494,22 @@ def flydsl_fp8_paged_mqa_logits_gfx950( kv_indices, max_model_len, *, + Preshuffle=True, KVBlockSize=KV_BLOCK_SIZE, SplitKV=None, TotalCuCount=None, stream=None, ): - """Run the H32/H64 D128/KVB64/preshuffle gfx950 mapping.""" + """Paged FP8 MQA logits (decode) for H in {32, 64}, next_n in {1, 2}, KVB=64. + + Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract on + the gfx950 16x16x128 mapping. Requires a preshuffled ``shuffle_weight(16,16)`` + KV cache. + """ if get_gfx() != _GFX950: raise RuntimeError(f"gfx950 kernel requested on {get_gfx()}") + if not Preshuffle: + raise ValueError("requires Preshuffle=True") batch_size, next_n, heads, head_dim = q_fp8.shape if heads not in SUPPORTED_HEADS or (head_dim, int(KVBlockSize)) != ( HEAD_DIM, @@ -371,7 +523,7 @@ def flydsl_fp8_paged_mqa_logits_gfx950( if kv_cache.dtype != torch.uint8: raise ValueError("kv_cache must contain preshuffled uint8 FP8 data") _, block_size, one, index_dim = kv_cache.shape - if block_size != KV_BLOCK_SIZE or one != 1 or index_dim != HEAD_DIM + 4: + if block_size != KV_BLOCK_SIZE or one != 1 or index_dim != INDEX_DIM: raise ValueError(f"unexpected KV cache shape {tuple(kv_cache.shape)}") context_lens = context_lens.reshape(batch_size) @@ -396,10 +548,8 @@ def flydsl_fp8_paged_mqa_logits_gfx950( else max(1, min(max_pages, int(SplitKV))) ) grid_blocks = batch_size * split_kv - launcher = _compile( - index_dim=int(index_dim), num_heads=int(heads), next_n=int(next_n) - ) - launcher.compile_hints = {**DEFAULT_COMPILE_HINTS, "waves_per_eu": 2} + launcher = _compile(num_heads=int(heads), next_n=int(next_n)) + launcher.compile_hints = dict(DEFAULT_COMPILE_HINTS) stream = stream or torch.cuda.current_stream(q_fp8.device) with torch.cuda.device(q_fp8.device.index): _run_compiled( @@ -421,3 +571,6 @@ def flydsl_fp8_paged_mqa_logits_gfx950( stream, ) return out_logits + + +flydsl_fp8_paged_mqa_logits_gfx950 = flydsl_fp8_paged_mqa_logits diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index a36e812fa4a..7811deaecc7 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -240,8 +240,6 @@ def _verify_paged_mqa_logits( split_kv, block_size, preshuffle, - variant, - chunk_k, ): _split_kv = None if split_kv == 0 else split_kv kv_cache_kernel, out = _kernel_inputs( @@ -273,8 +271,6 @@ def _verify_paged_mqa_logits( Preshuffle=preshuffle, KVBlockSize=block_size, SplitKV=_split_kv, - ChunkK=chunk_k, - variant=variant, ) got_mask = got == float("-inf") @@ -309,8 +305,6 @@ def _bench_flydsl_paged_kernel( preshuffle, block_size, split_kv, - chunk_k, - variant, ): _split_kv = None if split_kv == 0 else split_kv @@ -326,8 +320,6 @@ def fn(): Preshuffle=preshuffle, KVBlockSize=block_size, SplitKV=_split_kv, - ChunkK=chunk_k, - variant=variant, ) with torch.inference_mode(): @@ -343,10 +335,8 @@ def test_fp8_paged_mqa_logits( avg_kv_length, q_dtype, split_kv=0, - block_size=1, - preshuffle=False, - variant=None, - chunk_k=128, + block_size=64, + preshuffle=True, bench=False, ): inp = _build_inputs( @@ -366,16 +356,12 @@ def test_fp8_paged_mqa_logits( split_kv, block_size, preshuffle, - variant, - chunk_k, ) ret = { "gfx": get_gfx(), "kvb": block_size, "preshuffle": preshuffle, "split_kv": "auto" if split_kv == 0 else split_kv, - "variant": variant or "default", - "chunk_k": chunk_k, "flydsl err": err, } if not bench: @@ -393,41 +379,38 @@ def test_fp8_paged_mqa_logits( preshuffle=preshuffle, block_size=block_size, split_kv=split_kv, - chunk_k=chunk_k, - variant=variant, ) ) - if block_size == 1 and not preshuffle: - from aiter.ops.triton.attention.pa_mqa_logits import ( - deepgemm_fp8_paged_mqa_logits, - ) + from aiter.ops.triton.attention.pa_mqa_logits import ( + deepgemm_fp8_paged_mqa_logits, + ) + + out_ref = torch.full( + (batch_size * next_n, inp.max_model_len), + float("-inf"), + device="cuda", + dtype=torch.float32, + ) - out_ref = torch.full( - (batch_size * next_n, inp.max_model_len), - float("-inf"), - device="cuda", - dtype=torch.float32, + def fn_ref(): + return deepgemm_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out_ref, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + ChunkK=REF_CHUNK_K, + Preshuffle=True, + KVBlockSize=block_size, + WavePerEU=REF_WAVE_PER_EU, ) - def fn_ref(): - return deepgemm_fp8_paged_mqa_logits( - inp.q_fp8, - inp.kv_cache_fp8, - inp.weights, - out_ref, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - ChunkK=REF_CHUNK_K, - Preshuffle=False, - KVBlockSize=1, - WavePerEU=REF_WAVE_PER_EU, - ) - - with torch.inference_mode(): - _, ref_us = run_perftest(fn_ref) - ret["gluon us"] = float(ref_us) + with torch.inference_mode(): + _, ref_us = run_perftest(fn_ref) + ret["gluon us"] = float(ref_us) return ret @@ -440,7 +423,6 @@ def fn_ref(): "avg_kv_length": 1024, "q_dtype": "fnuz", } -_SHAPES = [(64, 64), (64, 128), (128, 64), (128, 128)] _AXES = ( "batch_size", "next_n", @@ -454,57 +436,36 @@ def fn_ref(): def _c(**kw): - return {**_BASE, "block_size": 64, **kw} + return {**_BASE, "block_size": 64, "preshuffle": True, **kw} def default_cases(): cases = [] - for heads, head_dim in _SHAPES: - for kvb in (1, 64): - cases.append(_c(heads=heads, head_dim=head_dim, block_size=kvb)) - for kvb in (16, 64): - cases.append( - _c(heads=heads, head_dim=head_dim, block_size=kvb, preshuffle=True) - ) + for heads in (32, 64): + for next_n in (1, 2): + cases.append(_c(heads=heads, next_n=next_n)) cases += [_c(batch_size=b, next_n=n) for b, n in ((1, 1), (1, 2), (4, 2), (8, 1))] cases += [_c(avg_kv_length=kv) for kv in (128, 8192)] cases += [_c(split_kv=sk) for sk in (1, 4)] cases.append(_c(avg_kv_length=128, split_kv=4)) cases.append(_c(avg_kv_length=8192, split_kv=1)) - cases.append(_c(split_kv=4, preshuffle=True)) - cases.append(_c(split_kv=1, block_size=16, preshuffle=True)) - cases += [_c(variant=v) for v in ("paged_w2", "paged_w4")] - cases.append(_c(chunk_k=64, variant="paged_w2")) - cases.append(_c(chunk_k=256, variant="paged_w4")) - cases.append(_c(chunk_k=256)) - cases.append(_c(preshuffle=True, variant="paged_w4")) return cases def exhaustive_cases(): prod = itertools.product( [(1, 1), (1, 2), (2, 1), (2, 2), (4, 2), (8, 1)], - [64, 128], - [64, 128], + [32, 64], + [128], [128, 1024, 8192], ["fnuz"], [0, 1, 4], - [1, 64], + [64], ) - cases = [ - dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, sk, kvb))) + return [ + {**dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, sk, kvb))), "preshuffle": True} for (bs, nn), nh, hd, kv, qd, sk, kvb in prod ] - return cases + [ - {**dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, 0, kvb))), "preshuffle": True} - for bs, nn, nh, hd, kv, qd, kvb in [ - (1, 1, 64, 128, 1024, "fnuz", 16), - (1, 2, 64, 128, 1024, "fnuz", 64), - (2, 1, 128, 128, 8192, "fnuz", 64), - (1, 1, 64, 64, 1024, "fnuz", 16), - (4, 2, 64, 128, 8192, "fnuz", 64), - ] - ] def gluon_ab_shapes(): @@ -538,9 +499,7 @@ def run_profile(args): args.head_dim, args.split_kv, args.kv_block_size, - args.preshuffle, - args.variant, - args.chunk_k, + True, ) _split_kv = None if args.split_kv == 0 else args.split_kv @@ -553,12 +512,9 @@ def run(): inp.context_lens, inp.block_tables, inp.max_model_len, - Preshuffle=args.preshuffle, + Preshuffle=True, KVBlockSize=args.kv_block_size, - ChunkK=args.chunk_k, SplitKV=_split_kv, - WavePerEU=args.wave_per_eu, - variant=args.variant, ) run() @@ -569,7 +525,7 @@ def run(): print( f"# gfx={get_gfx()} B={args.batch} nn={args.next_n} H={args.heads} " f"D={args.head_dim} kv_len={args.kv_len} kvb={args.kv_block_size} " - f"preshuffle={args.preshuffle} chunk_k={args.chunk_k} iters={args.iters}" + f"preshuffle=True iters={args.iters}" ) print(f"# KV requested: {kv_bytes / 1e6:.1f} MB") torch.cuda.synchronize() @@ -595,8 +551,8 @@ def run_gluon_ab(args): head_dim=D, avg_kv_length=kv_len, q_dtype="fnuz", - block_size=1, - preshuffle=False, + block_size=64, + preshuffle=True, bench=True, ) fly_us = ret["flydsl us"] @@ -642,9 +598,6 @@ def main(): action="store_true", help="run the full cartesian product instead of the curated matrix", ) - parser.add_argument( - "--no-preshuffle", action="store_true", help="skip preshuffle cases" - ) parser.add_argument( "--bench", action="store_true", @@ -653,7 +606,7 @@ def main(): parser.add_argument( "--compare-gluon", action="store_true", - help="run the production Gluon A/B decode shape sweep (KVBlockSize=1)", + help="run the Gluon A/B decode shape sweep (KVB=64, preshuffle)", ) parser.add_argument( "--profile", @@ -665,13 +618,9 @@ def main(): parser.add_argument("--heads", type=int, default=64) parser.add_argument("--head-dim", type=int, default=128) parser.add_argument("--kv-len", type=int, default=32768) - parser.add_argument("--kv-block-size", type=int, default=1) - parser.add_argument("--preshuffle", action="store_true") + parser.add_argument("--kv-block-size", type=int, default=64) parser.add_argument("--q-dtype", type=str, default="fnuz", choices=["fnuz", "fn"]) parser.add_argument("--split-kv", type=int, default=0) - parser.add_argument("--wave-per-eu", type=int, default=2) - parser.add_argument("--chunk-k", type=int, default=128) - parser.add_argument("--variant", type=str, default=None) parser.add_argument("--var-ratio", type=float, default=0.0) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--pool-blocks", type=int, default=0) @@ -681,8 +630,11 @@ def main(): args = parser.parse_args() if args.profile: - if args.preshuffle and args.kv_block_size % 16 != 0: - raise SystemExit("--preshuffle requires --kv-block-size divisible by 16") + if args.kv_block_size != 64 or args.heads not in (32, 64) or args.next_n not in ( + 1, + 2, + ): + raise SystemExit("kernel supports H in {32,64}, next_n in {1,2}, KVB=64") run_profile(args) return if args.compare_gluon: @@ -690,8 +642,6 @@ def main(): return cases = exhaustive_cases() if args.exhaustive else default_cases() - if args.no_preshuffle: - cases = [c for c in cases if not c.get("preshuffle")] df = [test_fp8_paged_mqa_logits(bench=args.bench, **c) for c in cases] df = pd.DataFrame(df) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index b76f05ec4e2..a457711bfaa 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -5,13 +5,11 @@ import argparse +import pytest import torch from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl import ( - flydsl_fp8_paged_mqa_logits, - flydsl_fp8_paged_mqa_logits_generic, -) +from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.test_common import run_perftest from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( @@ -30,6 +28,7 @@ HEAD_DIM = 128 KV_LEN = 32768 KV_BLOCK_SIZE = 64 +WIDE_MAX_MODEL_LEN = 1 << 20 def _inputs(heads, next_n=NEXT_N, kv_len=KV_LEN, batch=BATCH): @@ -97,6 +96,48 @@ def test_gfx950_nq1_and_short_pages(heads=32): _check(inp, kv_cache, out, f"H={heads} pages={pages}") +@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") +def test_gfx950_wide_output_no_tail_drop(): + """Row 512 crosses the signed i32 byte-offset boundary at this stride.""" + batch = 513 + inp, kv_cache, _ = _inputs(32, next_n=1, kv_len=KV_BLOCK_SIZE, batch=batch) + compact = torch.full( + (batch, inp.max_model_len), float("-inf"), dtype=torch.float32 + ) + wide = torch.full( + (batch, WIDE_MAX_MODEL_LEN), float("-inf"), dtype=torch.float32 + ) + + with torch.inference_mode(): + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache, + inp.weights, + compact, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + ) + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache, + inp.weights, + wide, + inp.context_lens, + inp.block_tables, + WIDE_MAX_MODEL_LEN, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + ) + torch.cuda.synchronize() + + valid_cols = KV_BLOCK_SIZE + assert torch.equal(wide[:, :valid_cols], compact[:, :valid_cols]) + assert torch.equal(wide[512:, :valid_cols], compact[512:, :valid_cols]) + + def _benchmark(heads): from aiter.ops.triton.attention.pa_mqa_logits import ( deepgemm_fp8_paged_mqa_logits, @@ -105,25 +146,11 @@ def _benchmark(heads): ) inp, kv_cache, out_new = _inputs(heads) - out_old = torch.full_like(out_new, float("-inf")) out_triton = torch.full_like(out_new, float("-inf")) def new_kernel(): return _launch(inp, kv_cache, out_new) - def old_kernel(): - return flydsl_fp8_paged_mqa_logits_generic( - inp.q_fp8, - kv_cache, - inp.weights, - out_old, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - Preshuffle=True, - KVBlockSize=KV_BLOCK_SIZE, - ) - def triton_kernel(): return deepgemm_fp8_paged_mqa_logits( inp.q_fp8, @@ -143,16 +170,13 @@ def triton_kernel(): print(f"H={heads} triton {triton_version} backend={backend}") with torch.inference_mode(): _, new_us = run_perftest(new_kernel, num_iters=50, num_warmup=8) - _, old_us = run_perftest(old_kernel, num_iters=50, num_warmup=8) _, triton_us = run_perftest(triton_kernel, num_iters=50, num_warmup=8) # WaveScope benchmarkPattern matches the first `time: {us} us` line (H=64). if heads == DEFAULT_HEADS: print(f"time: {new_us:.3f} us") else: print(f"time H={heads}: {new_us:.3f} us") - print(f"flydsl baseline: {old_us:.3f} us") print(f"{backend}: {triton_us:.3f} us") - print(f"vs flydsl: {old_us / new_us:.3f}x") print(f"vs {backend}: {triton_us / new_us:.3f}x") From 9a418792ba0b851f3cadde47dc8a4a10ec90427f Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 8 Sep 2026 13:14:09 +0000 Subject: [PATCH 36/52] style(flydsl): black-format gfx950 paged MQA logits and tests Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 19 +++++-------------- .../test_flydsl_fp8_paged_mqa_logits.py | 11 ++++++++--- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 8 ++------ 3 files changed, 15 insertions(+), 23 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 267269ac594..7c136b82ac1 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -95,9 +95,7 @@ def load_q_pack(q_i32, byte_base, lane_div_16): return _concat_i32x4(lo, hi) -def load_preshuffled_k_pack( - kv_i32, physical, tile_in_page, lane_mod_16, lane_div_16 -): +def load_preshuffled_k_pack(kv_i32, physical, tile_in_page, lane_mod_16, lane_div_16): """Load a shuffle_weight(16,16) K column as an i32x8 MFMA operand.""" base = ( physical * BLOCK_I32 @@ -223,9 +221,7 @@ def kernel( lane_mod_16 = umod(lane, MFMA_N) lane_div_16 = udiv(lane, MFMA_N) mma = fx.make_mma_atom( - fx.rocdl.cdna4.MFMA_Scale( - MFMA_M, MFMA_N, HEAD_DIM, fx.Float8E4M3FN - ) + fx.rocdl.cdna4.MFMA_Scale(MFMA_M, MFMA_N, HEAD_DIM, fx.Float8E4M3FN) ) q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) @@ -253,8 +249,7 @@ def kernel( [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(n_rows) ] w_rows = [ - [None for _ in range_constexpr(m_tiles)] - for _ in range_constexpr(n_rows) + [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(n_rows) ] for row in range_constexpr(n_rows): q_row = fx.Int32(row) @@ -313,9 +308,7 @@ def _store_page_bank(bank, b_slots, scale_slots): for slot in range_constexpr(B_RING): b_bank[slot].store(b_slots[slot]) scale_bank[slot].store( - fx.Vector.from_elements( - [scale_slots[slot].ir_value()], fx.Float32 - ) + fx.Vector.from_elements([scale_slots[slot].ir_value()], fx.Float32) ) def _load_page_bank(bank): @@ -359,9 +352,7 @@ def _compute_page(page_col, b_slots, scale_slots): ] for slot in range_constexpr(B_RING): scores = [ - mfma_scores( - mma, a_rows[row], b_slots[slot], m_tiles=m_tiles - ) + mfma_scores(mma, a_rows[row], b_slots[slot], m_tiles=m_tiles) for row in range_constexpr(n_rows) ] if slot > 0: diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 7811deaecc7..19a1c4b9ae4 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -630,9 +630,14 @@ def main(): args = parser.parse_args() if args.profile: - if args.kv_block_size != 64 or args.heads not in (32, 64) or args.next_n not in ( - 1, - 2, + if ( + args.kv_block_size != 64 + or args.heads not in (32, 64) + or args.next_n + not in ( + 1, + 2, + ) ): raise SystemExit("kernel supports H in {32,64}, next_n in {1,2}, KVB=64") run_profile(args) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index a457711bfaa..e39bca73cf5 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -101,12 +101,8 @@ def test_gfx950_wide_output_no_tail_drop(): """Row 512 crosses the signed i32 byte-offset boundary at this stride.""" batch = 513 inp, kv_cache, _ = _inputs(32, next_n=1, kv_len=KV_BLOCK_SIZE, batch=batch) - compact = torch.full( - (batch, inp.max_model_len), float("-inf"), dtype=torch.float32 - ) - wide = torch.full( - (batch, WIDE_MAX_MODEL_LEN), float("-inf"), dtype=torch.float32 - ) + compact = torch.full((batch, inp.max_model_len), float("-inf"), dtype=torch.float32) + wide = torch.full((batch, WIDE_MAX_MODEL_LEN), float("-inf"), dtype=torch.float32) with torch.inference_mode(): flydsl_fp8_paged_mqa_logits( From f5b390371d665baa2c47ee2174320b3e2fa2cf71 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 11 Sep 2026 09:26:50 +0000 Subject: [PATCH 37/52] dynamic next_n scheduling WIP. --- .../mqa_logits/fp8_paged_mqa_logits.py | 2 +- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 662 +++++++++--------- .../test_flydsl_fp8_paged_mqa_logits.py | 30 +- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 180 ++--- 4 files changed, 410 insertions(+), 464 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py index 374975f356a..3fc1c358894 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Paged FP8 MQA logits (decode) -- gfx950 16x16x128 mapping. +"""Paged FP8 MQA logits (decode) -- gfx950 ragged LDS mapping. The implementation lives in ``fp8_paged_mqa_logits_gfx950``. This module re-exports the public entry so existing import paths keep working. diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 7c136b82ac1..cc0deb54619 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""gfx950 H32/H64 D128/KVB64 preshuffled paged FP8 MQA-logits kernel. +"""gfx950 H32/D128/KVB64 ragged paged FP8 MQA-logits kernel. -Public decode path for ``flydsl_fp8_paged_mqa_logits``. Supports H in {32, 64}, -``next_n`` in {1, 2}, and KVBlockSize=64 with a preshuffled KV cache. +One 8-wave workgroup owns one ``(sequence, SplitKV slice)``. Waves 0..3 +compute the four N=16 tiles in one page and replay every live query row. +Waves 4..7 stream the next page directly into the idle LDS bank. """ from functools import lru_cache @@ -12,48 +13,53 @@ import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl.expr import arith, range_constexpr, rocdl +from flydsl.expr import arith, gpu, range_constexpr, rocdl from flydsl.expr.typing import T from aiter.jit.utils.chip_info import get_gfx from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from .. import buffer_ops +from ..mxfp4_gemm_common import _lds_ptr3 from ..tensor_shim import GTensor, _run_compiled, _to_raw Vec = fx.Vector -SUPPORTED_HEADS = (32, 64) +MAX_NN = 8 +NUM_HEADS = 32 HEAD_DIM = 128 INDEX_DIM = HEAD_DIM + 4 KV_BLOCK_SIZE = 64 MFMA_M = MFMA_N = 16 +MFMA_K = HEAD_DIM +M_TILES = NUM_HEADS // MFMA_M DREG = 4 B_RING = KV_BLOCK_SIZE // MFMA_N -BLOCK_I32 = KV_BLOCK_SIZE * INDEX_DIM // 4 -SCALE_I32 = KV_BLOCK_SIZE * HEAD_DIM // 4 TILE_I32 = HEAD_DIM * MFMA_N // 4 HALF_I32 = MFMA_N * 16 // 4 LANE_K_I32 = (HEAD_DIM // (64 // MFMA_N)) * MFMA_N // 4 -# Per physical page: 4 tiles × (2× dwordx4 K + 1× f32 scale). -PAGE_VMEM_LOADS = B_RING * 3 -_NEUTRAL_E8M0 = 0x7F7F7F7F - -_GFX950 = "gfx950" -DEFAULT_COMPILE_HINTS = { - "waves_per_eu": 2, - "fast_fp_math": True, -} +THREADS = 512 +WAVES = THREADS // 64 +COMPUTE_WAVES = 4 + +PAGE_BYTES = KV_BLOCK_SIZE * INDEX_DIM +PAGE_CHUNKS = PAGE_BYTES // 16 +K_BANKS = 2 +K_LDS_BYTES = K_BANKS * PAGE_BYTES +Q_ROW_BYTES = NUM_HEADS * HEAD_DIM +Q_LDS_BYTES = MAX_NN * Q_ROW_BYTES +W_ROW_BYTES = NUM_HEADS * 4 +W_LDS_BYTES = MAX_NN * W_ROW_BYTES +LDS_BYTES = K_LDS_BYTES + Q_LDS_BYTES + W_LDS_BYTES +Q_LDS_OFF = K_LDS_BYTES +W_LDS_OFF = Q_LDS_OFF + Q_LDS_BYTES -def wait_vmcnt(n): - """Explicit s_waitcnt vmcnt(n). Leaves n VM ops outstanding.""" - rocdl.sched_barrier(0) - rocdl.s_waitcnt(vmcnt=int(n)) - rocdl.sched_barrier(0) +_NEUTRAL_E8M0 = 0x7F7F7F7F +_GFX950 = "gfx950" +DEFAULT_COMPILE_HINTS = {"waves_per_eu": 2, "fast_fp_math": True} -# Via Uint32 because `//` on a signed Int32 lowers to floordivsi. def udiv(a, b): return fx.Int32(fx.Uint32(a) // fx.Uint32(b)) @@ -72,113 +78,105 @@ def imin(a, b): return (a <= b).select(a, b) -def guarded_store(pred, store_fn): - """Predicated side effect; masked buffer stores can fault on gfx950.""" +def _wait_all(): + rocdl.sched_barrier(0) + rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0) + rocdl.sched_barrier(0) + +def _guarded(pred, fn): @flyc.jit - def _guarded(_pred=pred, _store=store_fn): + def _do(_pred=pred, _fn=fn): if _pred: - _store() + _fn() - _guarded() + _do() -def _concat_i32x4(lo, hi): - return Vec(lo).shuffle(Vec(hi), list(range(DREG * 2))) +def _lds_f32x4(base_i32, byte_off): + ptr = fx.inttoptr( + fx.PointerType.get(T.f32, fx.AddressSpace.Shared, 16), + fx.Int32(base_i32 + byte_off), + ) + return fx.ptr_load(ptr, result_type=T.vec(4, T.f32)) -def load_q_pack(q_i32, byte_base, lane_div_16): - """Load one row's lane-owned contiguous K32 segment as a 256-bit operand.""" - off = (byte_base + lane_div_16 * 32) // fx.Int32(4) - lo = buffer_ops.buffer_load(q_i32.rsrc, off, vec_width=4, dtype=T.i32) - hi = buffer_ops.buffer_load(q_i32.rsrc, off + 4, vec_width=4, dtype=T.i32) - return _concat_i32x4(lo, hi) +def _lds_f32(base_i32, byte_off): + ptr = fx.inttoptr( + fx.PointerType.get(T.f32, fx.AddressSpace.Shared, 4), + fx.Int32(base_i32 + byte_off), + ) + return fx.Float32(fx.ptr_load(ptr)) -def load_preshuffled_k_pack(kv_i32, physical, tile_in_page, lane_mod_16, lane_div_16): - """Load a shuffle_weight(16,16) K column as an i32x8 MFMA operand.""" - base = ( - physical * BLOCK_I32 - + fx.Int32(tile_in_page * TILE_I32) - + lane_div_16 * LANE_K_I32 - + lane_mod_16 * 4 - ) - lo = buffer_ops.buffer_load( - kv_i32.rsrc, base, vec_width=4, dtype=T.i32, cache_modifier=2 - ) - hi = buffer_ops.buffer_load( - kv_i32.rsrc, - base + HALF_I32, - vec_width=4, - dtype=T.i32, - cache_modifier=2, +def _concat_i32x4(lo, hi): + return Vec(lo).shuffle(Vec(hi), list(range(8))) + + +def _lds_i32x4(base_i32, byte_off): + ptr = fx.inttoptr( + fx.PointerType.get(T.i32, fx.AddressSpace.Shared, 16), + fx.Int32(base_i32 + byte_off), ) + return fx.ptr_load(ptr, result_type=T.vec(4, T.i32)) + + +def _load_q_pack_lds(q_base, row, head, lane_div_16): + """Lane-owned contiguous K32 segment for a 16x16x128 A operand.""" + base = q_base + row * Q_ROW_BYTES + head * HEAD_DIM + lane_div_16 * 32 + return _concat_i32x4(_lds_i32x4(base, 0), _lds_i32x4(base, 16)) + + +def _load_k_pack_lds(k_base, tile, lane_mod_16, lane_div_16): + """Lane-owned B operand from one shuffle_weight(16,16) N=16 tile.""" + base_i32 = tile * TILE_I32 + lane_div_16 * LANE_K_I32 + lane_mod_16 * 4 + lo = _lds_i32x4(k_base, base_i32 * 4) + hi = _lds_i32x4(k_base, (base_i32 + HALF_I32) * 4) return _concat_i32x4(lo, hi) -def load_kv_scale(kv_i32, physical, token_in_page): - off = physical * BLOCK_I32 + SCALE_I32 + token_in_page - return fx.Float32( - buffer_ops.buffer_load( - kv_i32.rsrc, off, vec_width=1, dtype=T.f32, cache_modifier=2 +def _load_weight_frag_lds(w_base, row, mi, lane_div_16): + return Vec( + _lds_f32x4( + w_base, + row * W_ROW_BYTES + (mi * MFMA_M + lane_div_16 * DREG) * 4, ) ) -def mfma_scores(mma, a_tiles, b_pack, *, m_tiles): - """Return one 4-f32 score fragment per 16-head M tile.""" - neutral = fx.Int32(_NEUTRAL_E8M0) - b_frag = fx.make_rmem_tensor(DREG * 2, fx.Int32) - b_frag.store(b_pack) - scores = [] - for mi in range_constexpr(m_tiles): - a_frag = fx.make_rmem_tensor(DREG * 2, fx.Int32) - acc = fx.make_rmem_tensor(DREG, fx.Float32) - a_frag.store(a_tiles[mi]) - acc.store(Vec.filled(DREG, 0.0, fx.Float32)) - fx.gemm( - mma, - acc, - a_frag, - b_frag, - acc, - scale_a=neutral, - scale_b=neutral, +def _mfma_score(a_pack, b_pack): + result_ty = Vec.make_type(DREG, fx.Float32) + neutral = arith.constant(_NEUTRAL_E8M0, type=T.i32) + return Vec( + rocdl.mfma_scale_f32_16x16x128_f8f6f4( + result_ty, + [ + a_pack, + b_pack, + Vec.filled(DREG, 0.0, fx.Float32), + 0, + 0, + 0, + neutral, + 0, + neutral, + ], ) - scores.append(acc.load()) - return scores + ) -def reduce_scores(scores, weights, kv_scale, *, m_tiles): - """ReLU, weighted H reduction, positive KV scale, then wave reduction.""" +def _reduce_scores(scores, weights, kv_scale): zero = fx.Float32(0.0) total = zero - for mi in range_constexpr(m_tiles): - frag = Vec(scores[mi]) + for mi in range_constexpr(M_TILES): for ii in range_constexpr(DREG): - total = total + fx.Float32(frag[ii]).maximumf(zero) * weights[mi][ii] + total = total + fx.Float32(scores[mi][ii]).maximumf(zero) * weights[mi][ii] total = total * kv_scale total = total + total.shuffle_xor(16, 64) return total + total.shuffle_xor(32, 64) -def schedule_mfma_valu_pairs(*, m_tiles): - """Pair each 32-cycle MFMA with the prior tile's 12-op VALU fragment.""" - for _ in range_constexpr(m_tiles): - rocdl.sched_mfma(1) - rocdl.sched_group_barrier("valu", DREG * 3, 0) - - -@lru_cache(maxsize=8) -def device_cu_count(device_index: int) -> int: - try: - return torch.cuda.get_device_properties(device_index).multi_processor_count - except Exception: # noqa: BLE001 - return 256 - - def _make_out_row_view(logits, stride_out, row): - """Fold the row byte offset into the pointer using i64 arithmetic.""" byte = fx.Int64(fx.Uint32(row)) * fx.Int64(fx.Uint32(stride_out)) * 4 return GTensor( logits, @@ -189,16 +187,25 @@ def _make_out_row_view(logits, stride_out, row): ) -def _build_kernel(*, num_heads: int, next_n: int): - m_tiles = num_heads // MFMA_M - n_rows = int(next_n) - store_vmem = n_rows +@lru_cache(maxsize=8) +def device_cu_count(device_index: int) -> int: + try: + return torch.cuda.get_device_properties(device_index).multi_processor_count + except Exception: # noqa: BLE001 + return 256 + + +def _build_kernel(): kernel_name = ( - f"fp8_paged_mqa_logits_gfx950_H{num_heads}_D{HEAD_DIM}_" - f"bkv64_kvb{KV_BLOCK_SIZE}_r{B_RING}_nq{n_rows}_db2_ws1_pfp_sb0_nt_ps_flydsl" + "fp8_paged_mqa_logits_gfx950_H32_D128_bkv64_kvb64_" + "nn8_w8_lds2_mfma16_ntg2l_flydsl" ) - @flyc.kernel(name=kernel_name, known_block_size=[64, 1, 1]) + @fx.struct + class SharedStorage: + raw: fx.Array[fx.Uint8, LDS_BYTES, 16] + + @flyc.kernel(name=kernel_name, known_block_size=[THREADS, 1, 1]) def kernel( Q: fx.Tensor, KV_cache: fx.Tensor, @@ -206,226 +213,213 @@ def kernel( out_logits: fx.Tensor, context_lens: fx.Tensor, kv_indices: fx.Tensor, + next_n_lens: fx.Tensor, batch_size: fx.Int32, split_kv: fx.Int32, stride_q_batch: fx.Int32, - stride_q_next_n: fx.Int32, - stride_q_heads: fx.Int32, max_block_len: fx.Int32, stride_out: fx.Int32, ): - bid = fx.block_idx.x + tid = fx.Int32(fx.thread_idx.x) + wave = udiv(tid, 64) + lane = umod(tid, 64) + lane_div_16 = udiv(lane, 16) + lane_mod_16 = umod(lane, 16) + + bid = fx.Int32(fx.block_idx.x) pid_batch = umod(bid, batch_size) pid_split = udiv(bid, batch_size) - lane = umod(fx.thread_idx.x, 64) - lane_mod_16 = umod(lane, MFMA_N) - lane_div_16 = udiv(lane, MFMA_N) - mma = fx.make_mma_atom( - fx.rocdl.cdna4.MFMA_Scale(MFMA_M, MFMA_N, HEAD_DIM, fx.Float8E4M3FN) - ) - q_i32 = GTensor(Q, dtype=T.i32, shape=(-1,)) - kv_i32 = GTensor(KV_cache, dtype=T.i32, shape=(-1,)) - weight_t = GTensor(weights, dtype=T.f32, shape=(-1,)) + q_t = GTensor(Q, dtype=T.i32, shape=(-1,)) + kv_t = GTensor(KV_cache, dtype=T.i32, shape=(-1,)) + w_t = GTensor(weights, dtype=T.i32, shape=(-1,)) context_t = GTensor(context_lens, dtype=T.i32, shape=(-1,)) table_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) + nn_t = GTensor(next_n_lens, dtype=T.i32, shape=(-1,)) + nn = fx.Int32(nn_t[pid_batch]) context_len = fx.Int32(context_t[pid_batch]) page_count = uceildiv(context_len, fx.Int32(KV_BLOCK_SIZE)) pages_per_split = uceildiv(page_count, split_kv) page_lo = pid_split * pages_per_split page_hi = imin(page_lo + pages_per_split, page_count) - col_lo = page_lo * KV_BLOCK_SIZE - col_hi = page_hi * KV_BLOCK_SIZE - next_n_c = fx.Int32(n_rows) - out_rows = [ - _make_out_row_view( - out_logits, stride_out, pid_batch * next_n_c + fx.Int32(row) + lds = fx.SharedAllocator().allocate(SharedStorage).peek() + lds_base = fx.Int32(fx.ptrtoint(lds.raw.ptr)) + q_lds_base = lds_base + Q_LDS_OFF + w_lds_base = lds_base + W_LDS_OFF + + def _issue_dma(rsrc, dst_byte, src_byte, *, aux): + rocdl.raw_ptr_buffer_load_lds( + rsrc, + _lds_ptr3(lds_base, dst_byte), + fx.Int32(16), + fx.Int32(src_byte), + fx.Int32(0), + fx.Int32(0), + aux, ) - for row in range_constexpr(n_rows) - ] - a_rows = [ - [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(n_rows) - ] - w_rows = [ - [None for _ in range_constexpr(m_tiles)] for _ in range_constexpr(n_rows) - ] - for row in range_constexpr(n_rows): - q_row = fx.Int32(row) - q_base = pid_batch * stride_q_batch + q_row * stride_q_next_n - out_row = pid_batch * next_n_c + q_row - for mi in range_constexpr(m_tiles): - h = mi * MFMA_M + lane_mod_16 - byte_base = q_base + h * stride_q_heads - a_rows[row][mi] = load_q_pack(q_i32, byte_base, lane_div_16) - w_rows[row][mi] = fx.Vector( - buffer_ops.buffer_load( - weight_t.rsrc, - out_row * num_heads + mi * MFMA_M + lane_div_16 * DREG, - vec_width=DREG, - dtype=T.f32, - ) + + # Cached prologue: live Q rows and their f32 head weights. + q_batch_base = pid_batch * stride_q_batch + for rep in range_constexpr(Q_LDS_BYTES // (THREADS * 16)): + ci = tid + rep * THREADS + row = udiv(ci * 16, Q_ROW_BYTES) + pred = row < nn + + def _copy_q(_ci=ci): + byte = _ci * 16 + _issue_dma( + q_t.rsrc, + Q_LDS_OFF + byte, + q_batch_base + byte, + aux=0, ) - def _load_physical(page_col): - table_idx = pid_batch * max_block_len + udiv( - page_col, fx.Int32(KV_BLOCK_SIZE) + _guarded(pred, _copy_q) + + w_batch_base = pid_batch * (MAX_NN * W_ROW_BYTES) + ci_w = tid + w_row = udiv(ci_w * 16, W_ROW_BYTES) + + def _copy_w(_ci=ci_w): + byte = _ci * 16 + _issue_dma( + w_t.rsrc, + W_LDS_OFF + byte, + w_batch_base + byte, + aux=0, ) + + _guarded((ci_w < W_LDS_BYTES // 16) & (w_row < nn), _copy_w) + _wait_all() + gpu.barrier() + + def _load_physical(page): + safe_page = imin(page, page_hi - fx.Int32(1)) return fx.Int32( buffer_ops.buffer_load( - table_t.rsrc, table_idx, vec_width=1, dtype=T.i32, is_scalar=True + table_t.rsrc, + pid_batch * max_block_len + safe_page, + vec_width=1, + dtype=T.i32, + is_scalar=True, ) ) - def _issue_physical(physical): - b_slots, scale_slots = [], [] - for slot in range_constexpr(B_RING): - token = slot * MFMA_N + lane_mod_16 - b_slots.append( - load_preshuffled_k_pack( - kv_i32, - physical, - slot, - lane_mod_16, - lane_div_16, + def _copy_page(bank_base, page, copy_tid, copy_threads, reps): + physical = _load_physical(page) + for rep in range_constexpr(reps): + ci = copy_tid + rep * copy_threads + pred = ci < PAGE_CHUNKS + + def _copy(_chunk=ci, _physical=physical): + _issue_dma( + kv_t.rsrc, + bank_base + _chunk * 16, + _physical * PAGE_BYTES + _chunk * 16, + aux=2, ) - ) - scale_slots.append(load_kv_scale(kv_i32, physical, token)) - return b_slots, scale_slots - - def _make_page_bank(): - return ( - [ - fx.make_rmem_tensor(DREG * 2, fx.Int32) - for _ in range_constexpr(B_RING) - ], - [fx.make_rmem_tensor(1, fx.Float32) for _ in range_constexpr(B_RING)], - ) - - def _store_page_bank(bank, b_slots, scale_slots): - b_bank, scale_bank = bank - for slot in range_constexpr(B_RING): - b_bank[slot].store(b_slots[slot]) - scale_bank[slot].store( - fx.Vector.from_elements([scale_slots[slot].ir_value()], fx.Float32) - ) - - def _load_page_bank(bank): - b_bank, scale_bank = bank - return ( - [fx.Vector(b_bank[slot].load()) for slot in range_constexpr(B_RING)], - [ - fx.Float32(fx.Vector(scale_bank[slot].load())[0]) - for slot in range_constexpr(B_RING) - ], - ) - - def _make_physical_bank(): - return fx.make_rmem_tensor(1, fx.Int32) - - def _store_physical(bank, physical): - bank.store(fx.Vector.from_elements([physical.ir_value()], fx.Int32)) - def _load_physical_bank(bank): - return fx.Int32(fx.Vector(bank.load())[0]) + _guarded(pred, _copy) - def _store_page_row(row, page_col, logits): - q_row = fx.Int32(row) - q_limit = context_len - next_n_c + q_row - col = page_col + lane - logit = logits[0] - for slot in range_constexpr(1, B_RING): - logit = (lane_div_16 == slot).select(logits[slot], logit) - writer = (col < context_len) & (col <= q_limit) + def _store_logit(row, col, value): + out = _make_out_row_view(out_logits, stride_out, pid_batch * MAX_NN + row) - def _write(_out=out_rows[row], _col=col, _value=logit): + def _write(_out=out, _col=col, _value=value): _out[_col] = _value - guarded_store(writer, _write) + _guarded( + (lane_div_16 == 0) + & (col < context_len) + & (col <= context_len - nn + row), + _write, + ) - def _compute_page(page_col, b_slots, scale_slots): - prev_scores = None - prev_scale = None - page_logits = [ - [None for _ in range_constexpr(B_RING)] for _ in range_constexpr(n_rows) - ] - for slot in range_constexpr(B_RING): + def _compute_tile(bank_base, page, tile): + k_base = lds_base + bank_base + token_base = tile * MFMA_N + b_pack = _load_k_pack_lds(k_base, tile, lane_mod_16, lane_div_16) + scale = _lds_f32( + k_base, + KV_BLOCK_SIZE * HEAD_DIM + (token_base + lane_mod_16) * 4, + ) + row = fx.Int32(0) + while row < nn: + a_packs = [ + _load_q_pack_lds( + q_lds_base, + row, + mi * MFMA_M + lane_mod_16, + lane_div_16, + ) + for mi in range_constexpr(M_TILES) + ] + weights_frag = [ + _load_weight_frag_lds(w_lds_base, row, mi, lane_div_16) + for mi in range_constexpr(M_TILES) + ] scores = [ - mfma_scores(mma, a_rows[row], b_slots[slot], m_tiles=m_tiles) - for row in range_constexpr(n_rows) + _mfma_score(a_packs[mi], b_pack) for mi in range_constexpr(M_TILES) ] - if slot > 0: - schedule_mfma_valu_pairs(m_tiles=m_tiles) - for row in range_constexpr(n_rows): - page_logits[row][slot - 1] = reduce_scores( - prev_scores[row], - w_rows[row], - prev_scale, - m_tiles=m_tiles, - ) - prev_scores = scores - prev_scale = scale_slots[slot] - schedule_mfma_valu_pairs(m_tiles=m_tiles) - for row in range_constexpr(n_rows): - page_logits[row][B_RING - 1] = reduce_scores( - prev_scores[row], - w_rows[row], - prev_scale, - m_tiles=m_tiles, + value = _reduce_scores(scores, weights_frag, scale) + _store_logit( + row, + page * KV_BLOCK_SIZE + token_base + lane_mod_16, + value, ) - _store_page_row(row, page_col, page_logits[row]) + row = row + fx.Int32(1) if page_lo < page_hi: - current_bank = _make_page_bank() - next_bank = _make_page_bank() - b_init, scale_init = _issue_physical(_load_physical(col_lo)) - _store_page_bank(current_bank, b_init, scale_init) - current_physical = _make_physical_bank() - next_physical = _make_physical_bank() - if col_lo + fx.Int32(KV_BLOCK_SIZE) < col_hi: - _store_physical( - next_physical, - _load_physical(col_lo + fx.Int32(KV_BLOCK_SIZE)), - ) - page_col = col_lo - while page_col + fx.Int32(KV_BLOCK_SIZE) < col_hi: - next_b, next_scale = _issue_physical(_load_physical_bank(next_physical)) - _store_page_bank(next_bank, next_b, next_scale) - has_following = page_col + fx.Int32(2 * KV_BLOCK_SIZE) < col_hi - if has_following: - _store_physical( - current_physical, - _load_physical(page_col + fx.Int32(2 * KV_BLOCK_SIZE)), - ) - # Leave the next-page gather outstanding while consuming current. - wait_vmcnt(PAGE_VMEM_LOADS) - current_b, current_scale = _load_page_bank(current_bank) - _compute_page(page_col, current_b, current_scale) - - if has_following: - following_b, following_scale = _issue_physical( - _load_physical_bank(current_physical) + # All 512 threads fill the first page (two 16-byte issues max). + _copy_page( + fx.Int32(0), + page_lo, + tid, + fx.Int32(THREADS), + 2, + ) + _wait_all() + gpu.barrier() + + init_state = [fx.Int32(0)] + for page, state in range(page_lo, page_hi, fx.Int32(1), init=init_state): + page_i = fx.Int32(page) + bank = state[0] + bank_base = bank * PAGE_BYTES + next_bank = fx.Int32(1) - bank + next_bank_base = next_bank * PAGE_BYTES + + compute_pred = wave < COMPUTE_WAVES + + def _compute(_bank_base=bank_base, _page=page_i): + _compute_tile(_bank_base, _page, wave) + + _guarded(compute_pred, _compute) + + helper_tid = tid - COMPUTE_WAVES * 64 + next_page = page_i + fx.Int32(1) + + def _prefetch( + _bank_base=next_bank_base, + _page=next_page, + _helper_tid=helper_tid, + ): + _copy_page( + _bank_base, + _page, + _helper_tid, + fx.Int32((WAVES - COMPUTE_WAVES) * 64), + 3, ) - _store_page_bank(current_bank, following_b, following_scale) - if page_col + fx.Int32(3 * KV_BLOCK_SIZE) < col_hi: - _store_physical( - next_physical, - _load_physical(page_col + fx.Int32(3 * KV_BLOCK_SIZE)), - ) - # following-page loads plus current-page stores can remain. - wait_vmcnt(PAGE_VMEM_LOADS + store_vmem) - else: - # Only current-page stores may remain. - wait_vmcnt(store_vmem) - next_b, next_scale = _load_page_bank(next_bank) - _compute_page(page_col + fx.Int32(KV_BLOCK_SIZE), next_b, next_scale) - page_col = page_col + fx.Int32(2 * KV_BLOCK_SIZE) - - if page_col < col_hi: - wait_vmcnt(0) - current_b, current_scale = _load_page_bank(current_bank) - _compute_page(page_col, current_b, current_scale) + + _guarded( + (wave >= COMPUTE_WAVES) & (next_page < page_hi), + _prefetch, + ) + _wait_all() + gpu.barrier() + yield [next_bank] @flyc.jit def launch( @@ -435,12 +429,11 @@ def launch( out_logits, context_lens, kv_indices, + next_n_lens, grid_blocks, batch_size, split_kv, stride_q_batch, - stride_q_next_n, - stride_q_heads, max_block_len, stride_out, stream: fx.Stream = fx.Stream(None), # noqa: B008 @@ -453,16 +446,15 @@ def launch( out_logits, context_lens, kv_indices, + next_n_lens, batch_size, split_kv, stride_q_batch, - stride_q_next_n, - stride_q_heads, max_block_len, stride_out, ).launch( grid=(arith.index_cast(T.index, _to_raw(grid_blocks)), 1, 1), - block=(64, 1, 1), + block=(THREADS, 1, 1), stream=stream, ) @@ -471,9 +463,9 @@ def launch( return launch -@lru_cache(maxsize=8) -def _compile(*, num_heads: int, next_n: int): - return _build_kernel(num_heads=num_heads, next_n=next_n) +@lru_cache(maxsize=1) +def _compile(): + return _build_kernel() def flydsl_fp8_paged_mqa_logits( @@ -485,61 +477,78 @@ def flydsl_fp8_paged_mqa_logits( kv_indices, max_model_len, *, + next_n_lens=None, Preshuffle=True, KVBlockSize=KV_BLOCK_SIZE, SplitKV=None, TotalCuCount=None, stream=None, ): - """Paged FP8 MQA logits (decode) for H in {32, 64}, next_n in {1, 2}, KVB=64. - - Drop-in for the Triton ``deepgemm_fp8_paged_mqa_logits`` tensor contract on - the gfx950 16x16x128 mapping. Requires a preshuffled ``shuffle_weight(16,16)`` - KV cache. - """ + """Ragged paged FP8 MQA logits for padded Q ``[B, 8, 32, 128]``.""" if get_gfx() != _GFX950: raise RuntimeError(f"gfx950 kernel requested on {get_gfx()}") if not Preshuffle: raise ValueError("requires Preshuffle=True") - batch_size, next_n, heads, head_dim = q_fp8.shape - if heads not in SUPPORTED_HEADS or (head_dim, int(KVBlockSize)) != ( + if q_fp8.ndim != 4: + raise ValueError(f"q_fp8 must be rank 4, got shape {tuple(q_fp8.shape)}") + batch_size, max_nn, heads, head_dim = q_fp8.shape + if (max_nn, heads, head_dim, int(KVBlockSize)) != ( + MAX_NN, + NUM_HEADS, HEAD_DIM, KV_BLOCK_SIZE, ): - raise ValueError("requires H in {32, 64}, D=128, KVBlockSize=64") - if next_n not in (1, 2): - raise ValueError("requires next_n in {1, 2}") + raise ValueError("requires Q shape [B, 8, 32, 128] and KVBlockSize=64") + if not q_fp8.is_contiguous(): + raise ValueError("q_fp8 must be contiguous") if q_fp8.dtype != get_fp8_e4m3_dtype(): raise ValueError(f"q_fp8 must be native FP8 E4M3, got {q_fp8.dtype}") - if kv_cache.dtype != torch.uint8: - raise ValueError("kv_cache must contain preshuffled uint8 FP8 data") + if kv_cache.dtype != torch.uint8 or not kv_cache.is_contiguous(): + raise ValueError("kv_cache must be contiguous preshuffled uint8 FP8 data") _, block_size, one, index_dim = kv_cache.shape - if block_size != KV_BLOCK_SIZE or one != 1 or index_dim != INDEX_DIM: + if (block_size, one, index_dim) != (KV_BLOCK_SIZE, 1, INDEX_DIM): raise ValueError(f"unexpected KV cache shape {tuple(kv_cache.shape)}") + if weights.shape != (batch_size * MAX_NN, NUM_HEADS): + raise ValueError(f"weights must have shape {(batch_size * MAX_NN, NUM_HEADS)}") + if not weights.is_contiguous() or weights.dtype != torch.float32: + raise ValueError("weights must be contiguous float32") context_lens = context_lens.reshape(batch_size) - weights = weights.reshape(batch_size * next_n, heads) max_block_len = kv_indices.shape[-1] kv_indices = kv_indices.reshape(batch_size, max_block_len) - total_cu = ( - device_cu_count(q_fp8.device.index) - if TotalCuCount is None - else int(TotalCuCount) - ) - max_pages = max(1, (int(max_model_len) + KV_BLOCK_SIZE - 1) // KV_BLOCK_SIZE) - split_kv = ( - max( + if next_n_lens is None: + next_n_lens = torch.full( + (batch_size,), + MAX_NN, + dtype=torch.int32, + device=q_fp8.device, + ) + next_n_lens = next_n_lens.reshape(batch_size) + if next_n_lens.dtype != torch.int32 or next_n_lens.device != q_fp8.device: + raise ValueError("next_n_lens must be int32 on the same device as q_fp8") + + # One page fills all four compute SIMDs. Target three resident WGs/CU, + # matching the 50 KiB LDS allocation. Auto mode reads the runtime context + # maximum once; explicit SplitKV is graph-capture safe. + if SplitKV is None or int(SplitKV) <= 0: + real_pages = max( 1, - min( - max_pages, - (total_cu * 4 * 2 + batch_size - 1) // batch_size, - ), + (int(context_lens.max().item()) + KV_BLOCK_SIZE - 1) // KV_BLOCK_SIZE, ) - if SplitKV is None or int(SplitKV) <= 0 - else max(1, min(max_pages, int(SplitKV))) - ) + total_cu = ( + device_cu_count(q_fp8.device.index) + if TotalCuCount is None + else int(TotalCuCount) + ) + split_kv = max( + 1, + min(real_pages, (total_cu * 3 + batch_size - 1) // batch_size), + ) + else: + split_kv = max(1, int(SplitKV)) + grid_blocks = batch_size * split_kv - launcher = _compile(num_heads=int(heads), next_n=int(next_n)) + launcher = _compile() launcher.compile_hints = dict(DEFAULT_COMPILE_HINTS) stream = stream or torch.cuda.current_stream(q_fp8.device) with torch.cuda.device(q_fp8.device.index): @@ -551,12 +560,11 @@ def flydsl_fp8_paged_mqa_logits( out_logits, context_lens, kv_indices, + next_n_lens, int(grid_blocks), int(batch_size), int(split_kv), int(q_fp8.stride(0)), - int(q_fp8.stride(1)), - int(q_fp8.stride(2)), int(max_block_len), int(out_logits.stride(0)), stream, diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 19a1c4b9ae4..89398aa3019 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -417,8 +417,8 @@ def fn_ref(): _BASE = { "batch_size": 2, - "next_n": 2, - "heads": 64, + "next_n": 8, + "heads": 32, "head_dim": 128, "avg_kv_length": 1024, "q_dtype": "fnuz", @@ -440,11 +440,7 @@ def _c(**kw): def default_cases(): - cases = [] - for heads in (32, 64): - for next_n in (1, 2): - cases.append(_c(heads=heads, next_n=next_n)) - cases += [_c(batch_size=b, next_n=n) for b, n in ((1, 1), (1, 2), (4, 2), (8, 1))] + cases = [_c(batch_size=b) for b in (1, 2, 4, 8)] cases += [_c(avg_kv_length=kv) for kv in (128, 8192)] cases += [_c(split_kv=sk) for sk in (1, 4)] cases.append(_c(avg_kv_length=128, split_kv=4)) @@ -454,8 +450,8 @@ def default_cases(): def exhaustive_cases(): prod = itertools.product( - [(1, 1), (1, 2), (2, 1), (2, 2), (4, 2), (8, 1)], - [32, 64], + [(1, 8), (2, 8), (4, 8), (8, 8)], + [32], [128], [128, 1024, 8192], ["fnuz"], @@ -614,8 +610,8 @@ def main(): help="run a single-config timed loop (for rocprofv3 / isolation profiling)", ) parser.add_argument("--batch", type=int, default=16) - parser.add_argument("--next-n", type=int, default=2) - parser.add_argument("--heads", type=int, default=64) + parser.add_argument("--next-n", type=int, default=8) + parser.add_argument("--heads", type=int, default=32) parser.add_argument("--head-dim", type=int, default=128) parser.add_argument("--kv-len", type=int, default=32768) parser.add_argument("--kv-block-size", type=int, default=64) @@ -630,16 +626,8 @@ def main(): args = parser.parse_args() if args.profile: - if ( - args.kv_block_size != 64 - or args.heads not in (32, 64) - or args.next_n - not in ( - 1, - 2, - ) - ): - raise SystemExit("kernel supports H in {32,64}, next_n in {1,2}, KVB=64") + if args.kv_block_size != 64 or args.heads != 32 or args.next_n != 8: + raise SystemExit("kernel requires H=32, padded next_n=8, KVB=64") run_profile(args) return if args.compare_gluon: diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index e39bca73cf5..a971c65ec7b 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Focused B16/Nq2/H{32,64}/D128/KVB64 test for the gfx950 indexer mapping.""" +"""Focused long-context tests for the gfx950 padded ragged LDS indexer.""" import argparse @@ -12,40 +12,44 @@ from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from aiter.test_common import run_perftest +from op_tests.flydsl_tests.ragged_nn import ( + MAX_NN, + ref_padded_ragged, + sample_next_n_lens, +) from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( _build_inputs, _kernel_inputs, calc_diff, - ref_fp8_paged_mqa_logits, ) torch.set_default_device("cuda") BATCH = 16 -NEXT_N = 2 -HEADS = (32, 64) -DEFAULT_HEADS = 64 +HEADS = 32 HEAD_DIM = 128 KV_LEN = 32768 KV_BLOCK_SIZE = 64 WIDE_MAX_MODEL_LEN = 1 << 20 -def _inputs(heads, next_n=NEXT_N, kv_len=KV_LEN, batch=BATCH): +def _inputs(kv_len=KV_LEN, batch=BATCH, seed=1079): + torch.manual_seed(seed) inp = _build_inputs( batch, - next_n, - heads, + MAX_NN, + HEADS, HEAD_DIM, kv_len, get_fp8_e4m3_dtype(), block_size=KV_BLOCK_SIZE, ) - kv_cache, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM, True, KV_BLOCK_SIZE) - return inp, kv_cache, out + kv_cache, out = _kernel_inputs(inp, batch, MAX_NN, HEAD_DIM, True, KV_BLOCK_SIZE) + next_n_lens = sample_next_n_lens(batch, MAX_NN, seed=seed).cuda() + return inp, kv_cache, out, next_n_lens -def _launch(inp, kv_cache, out): +def _launch(inp, kv_cache, out, next_n_lens, *, max_model_len=None, split_kv=None): return flydsl_fp8_paged_mqa_logits( inp.q_fp8, kv_cache, @@ -53,79 +57,73 @@ def _launch(inp, kv_cache, out): out, inp.context_lens, inp.block_tables, - inp.max_model_len, + inp.max_model_len if max_model_len is None else max_model_len, + next_n_lens=next_n_lens, Preshuffle=True, KVBlockSize=KV_BLOCK_SIZE, + SplitKV=split_kv, ) -def _check(inp, kv_cache, out, tag): +def _check(inp, kv_cache, out, next_n_lens, tag): with torch.inference_mode(): - ref = ref_fp8_paged_mqa_logits( + ref = ref_padded_ragged( inp.q, inp.kv_cache_fp8, inp.weights, inp.context_lens, inp.block_tables, + next_n_lens, inp.max_model_len, inp.fp8_dtype, + max_nn=MAX_NN, block_size=KV_BLOCK_SIZE, ) - got = _launch(inp, kv_cache, out) + got = _launch(inp, kv_cache, out, next_n_lens) ref_mask = ref == float("-inf") got_mask = got == float("-inf") assert torch.equal(got_mask, ref_mask), f"{tag}: causal/padding -inf mask mismatch" diff = calc_diff(got.masked_fill(got_mask, 0), ref.masked_fill(ref_mask, 0)) assert diff < 1e-3, f"{tag} calc_diff={diff}" - print(f"correctness {tag}: pass calc_diff={float(diff):.3e}") -def test_gfx950_indexer_mapping(heads=DEFAULT_HEADS): - assert get_gfx() == "gfx950" - inp, kv_cache, out = _inputs(heads) - _check(inp, kv_cache, out, f"H={heads}") +@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") +def test_gfx950_ragged_long_context(): + inp, kv_cache, out, next_n_lens = _inputs() + _check(inp, kv_cache, out, next_n_lens, "ragged long context") -def test_gfx950_nq1_and_short_pages(heads=32): - assert get_gfx() == "gfx950" - inp, kv_cache, out = _inputs(heads, next_n=1) - _check(inp, kv_cache, out, f"H={heads} Nq=1") - for pages in (2, 3, 4): - inp, kv_cache, out = _inputs(heads, kv_len=pages * KV_BLOCK_SIZE, batch=2) - _check(inp, kv_cache, out, f"H={heads} pages={pages}") +@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") +@pytest.mark.parametrize("pages", [2, 3, 4, 7]) +def test_gfx950_ragged_short_pages(pages): + inp, kv_cache, out, next_n_lens = _inputs( + kv_len=pages * KV_BLOCK_SIZE, batch=2, seed=pages + ) + _check(inp, kv_cache, out, next_n_lens, f"pages={pages}") @pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") def test_gfx950_wide_output_no_tail_drop(): - """Row 512 crosses the signed i32 byte-offset boundary at this stride.""" - batch = 513 - inp, kv_cache, _ = _inputs(32, next_n=1, kv_len=KV_BLOCK_SIZE, batch=batch) - compact = torch.full((batch, inp.max_model_len), float("-inf"), dtype=torch.float32) - wide = torch.full((batch, WIDE_MAX_MODEL_LEN), float("-inf"), dtype=torch.float32) + """Padded output row 512 crosses the signed i32 byte-offset boundary.""" + batch = 65 + inp, kv_cache, compact, _ = _inputs(kv_len=KV_BLOCK_SIZE, batch=batch, seed=5) + next_n_lens = torch.ones(batch, dtype=torch.int32) + wide = torch.full( + (batch * MAX_NN, WIDE_MAX_MODEL_LEN), + float("-inf"), + dtype=torch.float32, + ) with torch.inference_mode(): - flydsl_fp8_paged_mqa_logits( - inp.q_fp8, + _launch(inp, kv_cache, compact, next_n_lens, split_kv=1) + _launch( + inp, kv_cache, - inp.weights, - compact, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - Preshuffle=True, - KVBlockSize=KV_BLOCK_SIZE, - ) - flydsl_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache, - inp.weights, wide, - inp.context_lens, - inp.block_tables, - WIDE_MAX_MODEL_LEN, - Preshuffle=True, - KVBlockSize=KV_BLOCK_SIZE, + next_n_lens, + max_model_len=WIDE_MAX_MODEL_LEN, + split_kv=1, ) torch.cuda.synchronize() @@ -134,87 +132,39 @@ def test_gfx950_wide_output_no_tail_drop(): assert torch.equal(wide[512:, :valid_cols], compact[512:, :valid_cols]) -def _benchmark(heads): - from aiter.ops.triton.attention.pa_mqa_logits import ( - deepgemm_fp8_paged_mqa_logits, - enable_gluon_pa_mqa_logits, - triton_version, - ) - - inp, kv_cache, out_new = _inputs(heads) - out_triton = torch.full_like(out_new, float("-inf")) - - def new_kernel(): - return _launch(inp, kv_cache, out_new) +def _benchmark(): + inp, kv_cache, out, next_n_lens = _inputs() - def triton_kernel(): - return deepgemm_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache, - inp.weights, - out_triton, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - ChunkK=256, - Preshuffle=True, - KVBlockSize=KV_BLOCK_SIZE, - WavePerEU=2, - ) + def kernel(): + return _launch(inp, kv_cache, out, next_n_lens) - backend = "gluon" if enable_gluon_pa_mqa_logits else "triton-jit" - print(f"H={heads} triton {triton_version} backend={backend}") with torch.inference_mode(): - _, new_us = run_perftest(new_kernel, num_iters=50, num_warmup=8) - _, triton_us = run_perftest(triton_kernel, num_iters=50, num_warmup=8) - # WaveScope benchmarkPattern matches the first `time: {us} us` line (H=64). - if heads == DEFAULT_HEADS: - print(f"time: {new_us:.3f} us") - else: - print(f"time H={heads}: {new_us:.3f} us") - print(f"{backend}: {triton_us:.3f} us") - print(f"vs {backend}: {triton_us / new_us:.3f}x") + _, us = run_perftest(kernel, num_iters=50, num_warmup=8) + print(f"time: {us:.3f} us") -def _profile(heads): - inp, kv_cache, out = _inputs(heads) +def _profile(): + inp, kv_cache, out, next_n_lens = _inputs() with torch.inference_mode(): - _launch(inp, kv_cache, out) + _launch(inp, kv_cache, out, next_n_lens) torch.cuda.synchronize() - print(f"profile launch H={heads}: pass") + print("profile launch: pass") def main(): parser = argparse.ArgumentParser() parser.add_argument("--profile", action="store_true") parser.add_argument("--time", action="store_true") - parser.add_argument( - "--heads", - type=int, - choices=HEADS, - default=None, - help="Restrict to one head count (default: both; --time still prints H=64 first)", - ) args = parser.parse_args() - heads = (args.heads,) if args.heads is not None else HEADS if args.time: - # H=64 first so WaveScope still sees `time: {us} us` as the production line. - timed = tuple(h for h in (DEFAULT_HEADS,) + heads if h in heads) - seen = set() - ordered = [] - for h in timed: - if h not in seen: - ordered.append(h) - seen.add(h) - for h in ordered: - _benchmark(h) + _benchmark() elif args.profile: - for h in heads: - _profile(h) + _profile() else: - for h in heads: - test_gfx950_indexer_mapping(h) - test_gfx950_nq1_and_short_pages(h) + test_gfx950_ragged_long_context() + for pages in (2, 3, 4, 7): + test_gfx950_ragged_short_pages(pages) + test_gfx950_wide_output_no_tail_drop() if __name__ == "__main__": From d1cbe07b57b665826f0237a3cadd83eb29abe66c Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 11 Sep 2026 10:08:31 +0000 Subject: [PATCH 38/52] cleanup --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 28 +++++++------------ 1 file changed, 10 insertions(+), 18 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index cc0deb54619..c1ebc16382b 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -145,24 +145,16 @@ def _load_weight_frag_lds(w_base, row, mi, lane_div_16): def _mfma_score(a_pack, b_pack): - result_ty = Vec.make_type(DREG, fx.Float32) - neutral = arith.constant(_NEUTRAL_E8M0, type=T.i32) - return Vec( - rocdl.mfma_scale_f32_16x16x128_f8f6f4( - result_ty, - [ - a_pack, - b_pack, - Vec.filled(DREG, 0.0, fx.Float32), - 0, - 0, - 0, - neutral, - 0, - neutral, - ], - ) - ) + atom = fx.make_mma_atom(fx.rocdl.cdna4.MFMA_Scale(16, 16, 128, fx.Float8E4M3FN)) + a_frag = fx.make_rmem_tensor(8, fx.Int32) + a_frag.store(Vec(a_pack)) + b_frag = fx.make_rmem_tensor(8, fx.Int32) + b_frag.store(Vec(b_pack)) + c_frag = fx.make_rmem_tensor(DREG, fx.Float32) + c_frag.store(Vec.filled(DREG, 0.0, fx.Float32)) + ident = fx.Int32(_NEUTRAL_E8M0) + fx.gemm(atom, c_frag, a_frag, b_frag, c_frag, scale_a=ident, scale_b=ident) + return Vec(c_frag.load()) def _reduce_scores(scores, weights, kv_scale): From 30bb0022c54ba9e674285afdb63856b45fadacc5 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Fri, 11 Sep 2026 11:33:09 +0000 Subject: [PATCH 39/52] feat(flydsl): accept compact Q[B, next_n, 32, 128] for gfx950 paged MQA logits Match the Gluon contract so next_n can be 1..8 without padding Q to 8 rows; optional next_n_lens still selects live rows per sequence. Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 46 ++-- op_tests/flydsl_tests/ragged_nn.py | 99 ++++++++ ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 2 +- ...t_flydsl_fp8_paged_mqa_logits_ragged_nn.py | 212 ++++++++++++++++++ 4 files changed, 342 insertions(+), 17 deletions(-) create mode 100644 op_tests/flydsl_tests/ragged_nn.py create mode 100644 op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index c1ebc16382b..24b96d19c14 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -208,6 +208,8 @@ def kernel( next_n_lens: fx.Tensor, batch_size: fx.Int32, split_kv: fx.Int32, + rows_per_batch: fx.Int32, + has_next_n_lens: fx.Int32, stride_q_batch: fx.Int32, max_block_len: fx.Int32, stride_out: fx.Int32, @@ -229,7 +231,8 @@ def kernel( table_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) nn_t = GTensor(next_n_lens, dtype=T.i32, shape=(-1,)) - nn = fx.Int32(nn_t[pid_batch]) + ragged_nn = imin(fx.Int32(nn_t[pid_batch]), rows_per_batch) + nn = (has_next_n_lens != 0).select(ragged_nn, rows_per_batch) context_len = fx.Int32(context_t[pid_batch]) page_count = uceildiv(context_len, fx.Int32(KV_BLOCK_SIZE)) pages_per_split = uceildiv(page_count, split_kv) @@ -270,7 +273,7 @@ def _copy_q(_ci=ci): _guarded(pred, _copy_q) - w_batch_base = pid_batch * (MAX_NN * W_ROW_BYTES) + w_batch_base = pid_batch * (rows_per_batch * W_ROW_BYTES) ci_w = tid w_row = udiv(ci_w * 16, W_ROW_BYTES) @@ -316,7 +319,9 @@ def _copy(_chunk=ci, _physical=physical): _guarded(pred, _copy) def _store_logit(row, col, value): - out = _make_out_row_view(out_logits, stride_out, pid_batch * MAX_NN + row) + out = _make_out_row_view( + out_logits, stride_out, pid_batch * rows_per_batch + row + ) def _write(_out=out, _col=col, _value=value): _out[_col] = _value @@ -425,6 +430,8 @@ def launch( grid_blocks, batch_size, split_kv, + rows_per_batch, + has_next_n_lens, stride_q_batch, max_block_len, stride_out, @@ -441,6 +448,8 @@ def launch( next_n_lens, batch_size, split_kv, + rows_per_batch, + has_next_n_lens, stride_q_batch, max_block_len, stride_out, @@ -476,21 +485,26 @@ def flydsl_fp8_paged_mqa_logits( TotalCuCount=None, stream=None, ): - """Ragged paged FP8 MQA logits for padded Q ``[B, 8, 32, 128]``.""" + """Paged FP8 MQA logits for contiguous Q ``[B, next_n, 32, 128]``. + + ``next_n`` may be 1..8. The compiled kernel reserves LDS for eight rows; + ``next_n_lens`` optionally selects fewer live rows per sequence. + """ if get_gfx() != _GFX950: raise RuntimeError(f"gfx950 kernel requested on {get_gfx()}") if not Preshuffle: raise ValueError("requires Preshuffle=True") if q_fp8.ndim != 4: raise ValueError(f"q_fp8 must be rank 4, got shape {tuple(q_fp8.shape)}") - batch_size, max_nn, heads, head_dim = q_fp8.shape - if (max_nn, heads, head_dim, int(KVBlockSize)) != ( - MAX_NN, + batch_size, next_n, heads, head_dim = q_fp8.shape + if not 1 <= next_n <= MAX_NN: + raise ValueError(f"q_fp8 next_n must be in 1..{MAX_NN}, got {next_n}") + if (heads, head_dim, int(KVBlockSize)) != ( NUM_HEADS, HEAD_DIM, KV_BLOCK_SIZE, ): - raise ValueError("requires Q shape [B, 8, 32, 128] and KVBlockSize=64") + raise ValueError("requires Q shape [B, next_n, 32, 128] and KVBlockSize=64") if not q_fp8.is_contiguous(): raise ValueError("q_fp8 must be contiguous") if q_fp8.dtype != get_fp8_e4m3_dtype(): @@ -500,21 +514,19 @@ def flydsl_fp8_paged_mqa_logits( _, block_size, one, index_dim = kv_cache.shape if (block_size, one, index_dim) != (KV_BLOCK_SIZE, 1, INDEX_DIM): raise ValueError(f"unexpected KV cache shape {tuple(kv_cache.shape)}") - if weights.shape != (batch_size * MAX_NN, NUM_HEADS): - raise ValueError(f"weights must have shape {(batch_size * MAX_NN, NUM_HEADS)}") + if weights.shape != (batch_size * next_n, NUM_HEADS): + raise ValueError(f"weights must have shape {(batch_size * next_n, NUM_HEADS)}") if not weights.is_contiguous() or weights.dtype != torch.float32: raise ValueError("weights must be contiguous float32") + if out_logits.shape[0] != batch_size * next_n: + raise ValueError(f"out_logits first dimension must be {batch_size * next_n}") context_lens = context_lens.reshape(batch_size) max_block_len = kv_indices.shape[-1] kv_indices = kv_indices.reshape(batch_size, max_block_len) + has_next_n_lens = next_n_lens is not None if next_n_lens is None: - next_n_lens = torch.full( - (batch_size,), - MAX_NN, - dtype=torch.int32, - device=q_fp8.device, - ) + next_n_lens = context_lens next_n_lens = next_n_lens.reshape(batch_size) if next_n_lens.dtype != torch.int32 or next_n_lens.device != q_fp8.device: raise ValueError("next_n_lens must be int32 on the same device as q_fp8") @@ -556,6 +568,8 @@ def flydsl_fp8_paged_mqa_logits( int(grid_blocks), int(batch_size), int(split_kv), + int(next_n), + int(has_next_n_lens), int(q_fp8.stride(0)), int(max_block_len), int(out_logits.stride(0)), diff --git a/op_tests/flydsl_tests/ragged_nn.py b/op_tests/flydsl_tests/ragged_nn.py new file mode 100644 index 00000000000..21b77161529 --- /dev/null +++ b/op_tests/flydsl_tests/ragged_nn.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Compact or ragged-``next_n`` batch for paged FP8 MQA logits. + +Q is dense ``[B, next_n, H, D]``. For a uniform batch all rows are live, like +Gluon. For a ragged batch, ``next_n_lens[b]`` in ``{1..next_n}`` marks how many +leading rows are live. The gfx950 FlyDSL kernel consumes that optional vector; +unused Q/weight rows are ignored and unused logit rows stay ``-inf``. +""" + +from __future__ import annotations + +import torch + +from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( + ref_fp8_paged_mqa_logits, +) + +MAX_NN = 8 + + +def sample_next_n_lens(nq: int, max_nn: int = MAX_NN, seed: int = 1079) -> torch.Tensor: + g = torch.Generator(device="cpu") + g.manual_seed(seed + nq) + return torch.randint( + 1, + max_nn + 1, + (nq,), + generator=g, + dtype=torch.int32, + device="cpu", + ) + + +def padded_logits(batch: int, max_nn: int, max_model_len: int) -> torch.Tensor: + return torch.full( + (batch * max_nn, max_model_len), + float("-inf"), + dtype=torch.float32, + device="cuda", + ) + + +def scatter_seq_logits( + out: torch.Tensor, + seq: int, + max_nn: int, + row_logits: torch.Tensor, + n_rows: int, +) -> None: + """Copy ``n_rows`` of a compact ``[n, max_model_len]`` tensor into padded out.""" + dst = seq * max_nn + out[dst : dst + n_rows].copy_(row_logits[:n_rows]) + + +def ref_padded_ragged( + q_pad: torch.Tensor, + kv_cache_fp8: torch.Tensor, + weights_pad: torch.Tensor, + context_lens: torch.Tensor, + block_tables: torch.Tensor, + next_n_lens: torch.Tensor, + max_model_len: int, + fp8_dtype, + *, + max_nn: int = MAX_NN, + block_size: int = 64, +) -> torch.Tensor: + """Per-sequence compact ref, scattered into ``[B*max_nn, max_model_len]``.""" + batch = q_pad.shape[0] + out = padded_logits(batch, max_nn, max_model_len) + nn = next_n_lens.tolist() + for b, n in enumerate(nn): + n = int(n) + q_b = q_pad[b : b + 1, :n] + w_b = weights_pad[b * max_nn : b * max_nn + n] + compact = ref_fp8_paged_mqa_logits( + q_b, + kv_cache_fp8, + w_b, + context_lens[b : b + 1], + block_tables[b : b + 1], + max_model_len, + fp8_dtype, + block_size=block_size, + ) + scatter_seq_logits(out, b, max_nn, compact, n) + return out + + +def live_row_mask( + next_n_lens: torch.Tensor, max_nn: int, max_model_len: int +) -> torch.Tensor: + """True on live ``[B*max_nn, max_model_len]`` rows (all columns).""" + batch = next_n_lens.shape[0] + rows = torch.arange(max_nn, device=next_n_lens.device) + live = rows[None, :] < next_n_lens[:, None] + return live.reshape(batch * max_nn, 1).expand(-1, max_model_len) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index a971c65ec7b..239d37d30f7 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Focused long-context tests for the gfx950 padded ragged LDS indexer.""" +"""Focused long-context tests for the gfx950 compact/ragged LDS indexer.""" import argparse diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py new file mode 100644 index 00000000000..bb9e7ce8bdb --- /dev/null +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""Correctness of compact and ragged next_n paged MQA logits.""" + +import pytest +import torch + +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl.kernels.mqa_logits.fp8_paged_mqa_logits_gfx950 import ( + flydsl_fp8_paged_mqa_logits, +) +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from op_tests.flydsl_tests.ragged_nn import ( + MAX_NN, + live_row_mask, + padded_logits, + ref_padded_ragged, + sample_next_n_lens, +) +from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( + calc_diff, + kv_cache_cast_to_fp8, + preshuffle_kv_data, +) + +torch.set_default_device("cuda") + +HEADS = 32 +HEAD_DIM = 128 +KVB = 64 + + +def _batch(nq, ctx, max_nn=MAX_NN, seed=0): + torch.manual_seed(seed) + fp8 = get_fp8_e4m3_dtype() + pages = ctx // KVB + q = torch.randn((nq, max_nn, HEADS, HEAD_DIM), dtype=torch.bfloat16) + weights = torch.randn((nq * max_nn, HEADS), dtype=torch.float32) + context_lens = torch.full((nq,), ctx, dtype=torch.int32) + kv = torch.randn((max(nq * pages, 1), KVB, 1, HEAD_DIM), dtype=torch.bfloat16) + kv_fp8 = kv_cache_cast_to_fp8(kv, fp8) + ids = torch.arange(nq * pages, device="cuda", dtype=torch.int32) + tables = torch.zeros((nq, pages), dtype=torch.int32) + tables[:, :pages] = ids.reshape(nq, pages) + return q, weights, context_lens, tables, kv_fp8, fp8 + + +def _run_and_check( + q, + weights, + context_lens, + tables, + kv_fp8, + fp8, + next_n_lens, + ctx, + *, + split_kv=None, +): + max_nn = q.shape[1] + ref = ref_padded_ragged( + q, + kv_fp8, + weights, + context_lens, + tables, + next_n_lens, + ctx, + fp8, + max_nn=max_nn, + block_size=KVB, + ) + out = padded_logits(q.shape[0], max_nn, ctx) + got = flydsl_fp8_paged_mqa_logits( + q.to(fp8), + preshuffle_kv_data(kv_fp8, HEAD_DIM), + weights, + out, + context_lens, + tables, + ctx, + next_n_lens=next_n_lens, + SplitKV=split_kv, + ) + ref_inf = ref == float("-inf") + got_inf = got == float("-inf") + assert torch.equal(got_inf, ref_inf), "ragged -inf mask mismatch" + diff = calc_diff(got.masked_fill(got_inf, 0), ref.masked_fill(ref_inf, 0)) + assert diff < 1e-3, diff + return got, ref + + +@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") +@pytest.mark.parametrize("nn", [1, 2, 4, 8]) +def test_compact_uniform_next_n(nn): + nq, ctx, max_nn = 4, 256, nn + next_n_lens = torch.full((nq,), nn, dtype=torch.int32, device="cuda") + q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, max_nn) + ref = ref_padded_ragged( + q, + kv_fp8, + weights, + context_lens, + tables, + next_n_lens, + ctx, + fp8, + max_nn=max_nn, + block_size=KVB, + ) + out = padded_logits(nq, max_nn, ctx) + got = flydsl_fp8_paged_mqa_logits( + q.to(fp8), + preshuffle_kv_data(kv_fp8, HEAD_DIM), + weights, + out, + context_lens, + tables, + ctx, + ) + ref_inf = ref == float("-inf") + got_inf = got == float("-inf") + assert torch.equal(got_inf, ref_inf) + diff = calc_diff(got.masked_fill(got_inf, 0), ref.masked_fill(ref_inf, 0)) + assert diff < 1e-3, diff + + +@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") +@pytest.mark.parametrize("max_nn", [4, MAX_NN]) +def test_ragged_mix_unused_stay_neg_inf(max_nn): + nq, ctx = 8, 192 + next_n_lens = sample_next_n_lens(nq, max_nn, seed=1079).cuda() + assert int(next_n_lens.min()) >= 1 and int(next_n_lens.max()) <= max_nn + assert len(next_n_lens.unique()) >= 2 + q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, max_nn, seed=1079) + got, ref = _run_and_check( + q, + weights, + context_lens, + tables, + kv_fp8, + fp8, + next_n_lens, + ctx, + ) + mask = live_row_mask(next_n_lens, max_nn, ctx) + unused_rows = ~mask[:, 0] + assert unused_rows.any(), "fixture should include padded rows" + assert torch.all(ref[unused_rows] == float("-inf")) + assert torch.all(got[unused_rows] == float("-inf")) + + q2 = torch.randn_like(q) + for b, n in enumerate(next_n_lens.tolist()): + q2[b, :n].copy_(q[b, :n]) + ref2 = ref_padded_ragged( + q2, + kv_fp8, + weights, + context_lens, + tables, + next_n_lens, + ctx, + fp8, + max_nn=max_nn, + block_size=KVB, + ) + assert torch.equal(ref == float("-inf"), ref2 == float("-inf")) + inf = ref == float("-inf") + diff = calc_diff(ref.masked_fill(inf, 0), ref2.masked_fill(inf, 0)) + assert diff < 1e-3, diff + + +@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") +@pytest.mark.parametrize("nq", [1, 24]) +def test_nq_and_short_pages(nq): + ctx, max_nn = 128, MAX_NN + next_n_lens = sample_next_n_lens(nq, max_nn, seed=3).cuda() + q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, max_nn, seed=3) + got, ref = _run_and_check( + q, + weights, + context_lens, + tables, + kv_fp8, + fp8, + next_n_lens, + ctx, + ) + assert ref.shape == (nq * max_nn, ctx) + assert got.shape == ref.shape + empty = padded_logits(nq, max_nn, ctx) + mask = live_row_mask(next_n_lens, max_nn, ctx) + assert torch.equal(ref[~mask[:, 0]], empty[~mask[:, 0]]) + + +@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") +def test_odd_page_split_starts(): + nq, ctx = 3, 7 * KVB + next_n_lens = sample_next_n_lens(nq, MAX_NN, seed=19).cuda() + q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, seed=19) + _run_and_check( + q, + weights, + context_lens, + tables, + kv_fp8, + fp8, + next_n_lens, + ctx, + split_kv=3, + ) From 0c4b5b7ea34e0a71911f69591b1c822958687826 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 15 Sep 2026 09:40:17 +0000 Subject: [PATCH 40/52] perf(flydsl): raise wave priority around scaled MFMA Give compute waves s_setprio 3 during gemm so they win over helper G2L. Co-authored-by: Cursor --- .../flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 24b96d19c14..c8c5a9fa1eb 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -153,7 +153,9 @@ def _mfma_score(a_pack, b_pack): c_frag = fx.make_rmem_tensor(DREG, fx.Float32) c_frag.store(Vec.filled(DREG, 0.0, fx.Float32)) ident = fx.Int32(_NEUTRAL_E8M0) + rocdl.s_setprio(3) fx.gemm(atom, c_frag, a_frag, b_frag, c_frag, scale_a=ident, scale_b=ident) + rocdl.s_setprio(0) return Vec(c_frag.load()) From 1aaafcd257f2f8b3e913ae032db707d915c02d0a Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 15 Sep 2026 09:41:27 +0000 Subject: [PATCH 41/52] perf(flydsl): hold s_setprio across both scaled MFMAs per row Avoid dropping priority between the two 16x16x128 gems. Co-authored-by: Cursor --- .../flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index c8c5a9fa1eb..909536e8502 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -153,9 +153,7 @@ def _mfma_score(a_pack, b_pack): c_frag = fx.make_rmem_tensor(DREG, fx.Float32) c_frag.store(Vec.filled(DREG, 0.0, fx.Float32)) ident = fx.Int32(_NEUTRAL_E8M0) - rocdl.s_setprio(3) fx.gemm(atom, c_frag, a_frag, b_frag, c_frag, scale_a=ident, scale_b=ident) - rocdl.s_setprio(0) return Vec(c_frag.load()) @@ -358,9 +356,11 @@ def _compute_tile(bank_base, page, tile): _load_weight_frag_lds(w_lds_base, row, mi, lane_div_16) for mi in range_constexpr(M_TILES) ] + rocdl.s_setprio(3) scores = [ _mfma_score(a_packs[mi], b_pack) for mi in range_constexpr(M_TILES) ] + rocdl.s_setprio(0) value = _reduce_scores(scores, weights_frag, scale) _store_logit( row, From 2ab0067232b26469b53e91c23cdef32442430ceb Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 15 Sep 2026 09:42:05 +0000 Subject: [PATCH 42/52] perf(flydsl): hint waves_per_eu=6 to match LDS-bound occupancy Compiler was told 2 waves/EU while the CTA actually lands 6. Co-authored-by: Cursor --- .../flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 909536e8502..d488379f7c8 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -57,7 +57,7 @@ _NEUTRAL_E8M0 = 0x7F7F7F7F _GFX950 = "gfx950" -DEFAULT_COMPILE_HINTS = {"waves_per_eu": 2, "fast_fp_math": True} +DEFAULT_COMPILE_HINTS = {"waves_per_eu": 6, "fast_fp_math": True} def udiv(a, b): From d63d4bfb21bcd9482ba36f0000560892f622de12 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 15 Sep 2026 09:58:33 +0000 Subject: [PATCH 43/52] perf(flydsl): unroll next_n scoring by two rows Issue a pair of Q/weight LDS packs per iteration of the row loop. Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index d488379f7c8..44a25fda388 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -341,19 +341,18 @@ def _compute_tile(bank_base, page, tile): k_base, KV_BLOCK_SIZE * HEAD_DIM + (token_base + lane_mod_16) * 4, ) - row = fx.Int32(0) - while row < nn: + def _score_row(r): a_packs = [ _load_q_pack_lds( q_lds_base, - row, + r, mi * MFMA_M + lane_mod_16, lane_div_16, ) for mi in range_constexpr(M_TILES) ] weights_frag = [ - _load_weight_frag_lds(w_lds_base, row, mi, lane_div_16) + _load_weight_frag_lds(w_lds_base, r, mi, lane_div_16) for mi in range_constexpr(M_TILES) ] rocdl.s_setprio(3) @@ -363,10 +362,18 @@ def _compute_tile(bank_base, page, tile): rocdl.s_setprio(0) value = _reduce_scores(scores, weights_frag, scale) _store_logit( - row, + r, page * KV_BLOCK_SIZE + token_base + lane_mod_16, value, ) + + row = fx.Int32(0) + while row + fx.Int32(1) < nn: + _score_row(row) + _score_row(row + fx.Int32(1)) + row = row + fx.Int32(2) + while row < nn: + _score_row(row) row = row + fx.Int32(1) if page_lo < page_hi: From d2867b0669cb5b0b5182a332f5467f3b90476597 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 15 Sep 2026 09:59:17 +0000 Subject: [PATCH 44/52] perf(flydsl): unroll next_n scoring by four rows Give the compiler a larger Q/weight LDS pack window per tile. Co-authored-by: Cursor --- .../kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index 44a25fda388..dff14f7c1b4 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -368,6 +368,12 @@ def _score_row(r): ) row = fx.Int32(0) + while row + fx.Int32(3) < nn: + _score_row(row) + _score_row(row + fx.Int32(1)) + _score_row(row + fx.Int32(2)) + _score_row(row + fx.Int32(3)) + row = row + fx.Int32(4) while row + fx.Int32(1) < nn: _score_row(row) _score_row(row + fx.Int32(1)) From 4b2ac42114c71d47fa12042f03474ac514c9b50a Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 15 Sep 2026 09:59:57 +0000 Subject: [PATCH 45/52] perf(flydsl): fully unroll next_n scoring over MAX_NN Replace the 4/2/1 while ladder with predicated constexpr rows. Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 20 ++++++------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index dff14f7c1b4..c770bc2e178 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -367,20 +367,12 @@ def _score_row(r): value, ) - row = fx.Int32(0) - while row + fx.Int32(3) < nn: - _score_row(row) - _score_row(row + fx.Int32(1)) - _score_row(row + fx.Int32(2)) - _score_row(row + fx.Int32(3)) - row = row + fx.Int32(4) - while row + fx.Int32(1) < nn: - _score_row(row) - _score_row(row + fx.Int32(1)) - row = row + fx.Int32(2) - while row < nn: - _score_row(row) - row = row + fx.Int32(1) + for r in range_constexpr(MAX_NN): + + def _do(_r=r): + _score_row(fx.Int32(_r)) + + _guarded(fx.Int32(r) < nn, _do) if page_lo < page_hi: # All 512 threads fill the first page (two 16-byte issues max). From 6d6b90e1301a4dd05f3759730354c10e961a56cd Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 15 Sep 2026 13:33:39 +0000 Subject: [PATCH 46/52] perf(flydsl): cache paged MQA logit stores for the following topk Write logits with buffer cache policy 0 so L2 can retain the row for the next kernel; keep KV page DMA nontemporal so those loads do not evict it. Co-authored-by: Cursor --- .../kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index c770bc2e178..cd214b9ab51 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -174,7 +174,7 @@ def _make_out_row_view(logits, stride_out, row): logits, dtype=T.f32, shape=(-1,), - cache_modifier=2, + cache_modifier=0, static_bytes_offset_i64=byte, ) @@ -324,7 +324,9 @@ def _store_logit(row, col, value): ) def _write(_out=out, _col=col, _value=value): - _out[_col] = _value + # Cached write-allocate so the following topk can hit L2. + # KV page DMA stays NT (aux=2) and does not evict these lines. + buffer_ops.buffer_store(_value, _out.rsrc, _col, cache_modifier=0) _guarded( (lane_div_16 == 0) @@ -341,6 +343,7 @@ def _compute_tile(bank_base, page, tile): k_base, KV_BLOCK_SIZE * HEAD_DIM + (token_base + lane_mod_16) * 4, ) + def _score_row(r): a_packs = [ _load_q_pack_lds( From 43be646cc43058a3759e82f9048601d41a3ff90d Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Thu, 17 Sep 2026 12:46:14 +0000 Subject: [PATCH 47/52] refactor(flydsl): drop dead gfx950 paged-MQA helpers and restore scalar logit stores Scalar GTensor writes already lower to cached buffer_store; the explicit store path and unused arith/B_RING bits did not change ISA. Share MAX_NN from the kernel in tests. Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 26 +++++++------------ op_tests/flydsl_tests/ragged_nn.py | 3 +-- 2 files changed, 11 insertions(+), 18 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index cd214b9ab51..fec9e18b0de 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -13,7 +13,7 @@ import flydsl.compiler as flyc import flydsl.expr as fx import torch -from flydsl.expr import arith, gpu, range_constexpr, rocdl +from flydsl.expr import gpu, range_constexpr, rocdl from flydsl.expr.typing import T from aiter.jit.utils.chip_info import get_gfx @@ -21,7 +21,7 @@ from .. import buffer_ops from ..mxfp4_gemm_common import _lds_ptr3 -from ..tensor_shim import GTensor, _run_compiled, _to_raw +from ..tensor_shim import GTensor, _run_compiled Vec = fx.Vector @@ -34,7 +34,6 @@ MFMA_K = HEAD_DIM M_TILES = NUM_HEADS // MFMA_M DREG = 4 -B_RING = KV_BLOCK_SIZE // MFMA_N TILE_I32 = HEAD_DIM * MFMA_N // 4 HALF_I32 = MFMA_N * 16 // 4 LANE_K_I32 = (HEAD_DIM // (64 // MFMA_N)) * MFMA_N // 4 @@ -73,11 +72,6 @@ def uceildiv(a, b): return fx.Int32((fx.Uint32(a) + fx.Uint32(b) - 1) // fx.Uint32(b)) -def imin(a, b): - a, b = fx.Int32(a), fx.Int32(b) - return (a <= b).select(a, b) - - def _wait_all(): rocdl.sched_barrier(0) rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0) @@ -174,7 +168,6 @@ def _make_out_row_view(logits, stride_out, row): logits, dtype=T.f32, shape=(-1,), - cache_modifier=0, static_bytes_offset_i64=byte, ) @@ -231,13 +224,16 @@ def kernel( table_t = GTensor(kv_indices, dtype=T.i32, shape=(-1,)) nn_t = GTensor(next_n_lens, dtype=T.i32, shape=(-1,)) - ragged_nn = imin(fx.Int32(nn_t[pid_batch]), rows_per_batch) + def _imin(a, b): + return (a <= b).select(a, b) + + ragged_nn = _imin(fx.Int32(nn_t[pid_batch]), rows_per_batch) nn = (has_next_n_lens != 0).select(ragged_nn, rows_per_batch) context_len = fx.Int32(context_t[pid_batch]) page_count = uceildiv(context_len, fx.Int32(KV_BLOCK_SIZE)) pages_per_split = uceildiv(page_count, split_kv) page_lo = pid_split * pages_per_split - page_hi = imin(page_lo + pages_per_split, page_count) + page_hi = _imin(page_lo + pages_per_split, page_count) lds = fx.SharedAllocator().allocate(SharedStorage).peek() lds_base = fx.Int32(fx.ptrtoint(lds.raw.ptr)) @@ -291,7 +287,7 @@ def _copy_w(_ci=ci_w): gpu.barrier() def _load_physical(page): - safe_page = imin(page, page_hi - fx.Int32(1)) + safe_page = _imin(page, page_hi - fx.Int32(1)) return fx.Int32( buffer_ops.buffer_load( table_t.rsrc, @@ -324,9 +320,7 @@ def _store_logit(row, col, value): ) def _write(_out=out, _col=col, _value=value): - # Cached write-allocate so the following topk can hit L2. - # KV page DMA stays NT (aux=2) and does not evict these lines. - buffer_ops.buffer_store(_value, _out.rsrc, _col, cache_modifier=0) + _out[_col] = _value _guarded( (lane_div_16 == 0) @@ -464,7 +458,7 @@ def launch( max_block_len, stride_out, ).launch( - grid=(arith.index_cast(T.index, _to_raw(grid_blocks)), 1, 1), + grid=(fx.Int64(grid_blocks), 1, 1), block=(THREADS, 1, 1), stream=stream, ) diff --git a/op_tests/flydsl_tests/ragged_nn.py b/op_tests/flydsl_tests/ragged_nn.py index 21b77161529..c18be251894 100644 --- a/op_tests/flydsl_tests/ragged_nn.py +++ b/op_tests/flydsl_tests/ragged_nn.py @@ -13,12 +13,11 @@ import torch +from aiter.ops.flydsl.kernels.mqa_logits.fp8_paged_mqa_logits_gfx950 import MAX_NN from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( ref_fp8_paged_mqa_logits, ) -MAX_NN = 8 - def sample_next_n_lens(nq: int, max_nn: int = MAX_NN, seed: int = 1079) -> torch.Tensor: g = torch.Generator(device="cpu") From 705c287a1eb2b64c1fb12fdb5cb1f43317464b3e Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 22 Sep 2026 10:22:47 +0000 Subject: [PATCH 48/52] test(flydsl): keep paged MQA driver out of pytest collection Co-authored-by: Cursor --- op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index 89398aa3019..d7b54cf80b8 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -327,7 +327,7 @@ def fn(): return {"flydsl us": float(us)} -def test_fp8_paged_mqa_logits( +def run_case( batch_size, next_n, heads, @@ -540,7 +540,7 @@ def run(): def run_gluon_ab(args): rows = [] for B, nn, H, D, kv_len in gluon_ab_shapes(): - ret = test_fp8_paged_mqa_logits( + ret = run_case( batch_size=B, next_n=nn, heads=H, @@ -636,7 +636,7 @@ def main(): cases = exhaustive_cases() if args.exhaustive else default_cases() - df = [test_fp8_paged_mqa_logits(bench=args.bench, **c) for c in cases] + df = [run_case(bench=args.bench, **c) for c in cases] df = pd.DataFrame(df) try: summary = df.to_markdown(index=False) From 6295880b5b87cf85b15d4d6bc189f5796e10c76f Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 22 Sep 2026 11:07:44 +0000 Subject: [PATCH 49/52] cleanup(flydsl): remove stale paged MQA test paths Co-authored-by: Cursor --- .../mqa_logits/fp8_paged_mqa_logits_gfx950.py | 1 - .../test_flydsl_fp8_paged_mqa_logits.py | 14 ++++---------- .../test_flydsl_fp8_paged_mqa_logits_gfx950.py | 2 +- 3 files changed, 5 insertions(+), 12 deletions(-) diff --git a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py index fec9e18b0de..4bf0ec91081 100644 --- a/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -31,7 +31,6 @@ INDEX_DIM = HEAD_DIM + 4 KV_BLOCK_SIZE = 64 MFMA_M = MFMA_N = 16 -MFMA_K = HEAD_DIM M_TILES = NUM_HEADS // MFMA_M DREG = 4 TILE_I32 = HEAD_DIM * MFMA_N // 4 diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index d7b54cf80b8..e492b30abc3 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -217,12 +217,8 @@ def _build_inputs( ) -def _kernel_inputs(inp, batch_size, next_n, head_dim, preshuffle, block_size): - kv_cache_kernel = ( - preshuffle_kv_data(inp.kv_cache_fp8, head_dim) - if preshuffle - else inp.kv_cache_fp8 - ) +def _kernel_inputs(inp, batch_size, next_n, head_dim): + kv_cache_kernel = preshuffle_kv_data(inp.kv_cache_fp8, head_dim) out = torch.full( (batch_size * next_n, inp.max_model_len), float("-inf"), @@ -242,9 +238,7 @@ def _verify_paged_mqa_logits( preshuffle, ): _split_kv = None if split_kv == 0 else split_kv - kv_cache_kernel, out = _kernel_inputs( - inp, batch_size, next_n, head_dim, preshuffle, block_size - ) + kv_cache_kernel, out = _kernel_inputs(inp, batch_size, next_n, head_dim) with torch.inference_mode(): ref = ref_fp8_paged_mqa_logits( @@ -465,7 +459,7 @@ def exhaustive_cases(): def gluon_ab_shapes(): - H, D = 64, 128 + H, D = 32, 128 shapes = [] for B in (1, 4, 16, 64, 128): for avg_kv in (16384, 32768, 65536): diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index 239d37d30f7..b1d99e4b83a 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -44,7 +44,7 @@ def _inputs(kv_len=KV_LEN, batch=BATCH, seed=1079): get_fp8_e4m3_dtype(), block_size=KV_BLOCK_SIZE, ) - kv_cache, out = _kernel_inputs(inp, batch, MAX_NN, HEAD_DIM, True, KV_BLOCK_SIZE) + kv_cache, out = _kernel_inputs(inp, batch, MAX_NN, HEAD_DIM) next_n_lens = sample_next_n_lens(batch, MAX_NN, seed=seed).cuda() return inp, kv_cache, out, next_n_lens From 3075570a087caa3477c94a7796fe596242a6c0c3 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 22 Sep 2026 14:46:07 +0000 Subject: [PATCH 50/52] test(flydsl): put paged MQA logits on the aiter op_test template Move the gfx950 indexer tests to op_tests/test_*.py with @benchmark candidates, TFLOPS/TB/s, a torch-only reference, and one markdown table per signature so CI runs python3 op_tests/test_flydsl_fp8_paged_mqa_logits.py. Co-authored-by: Cursor --- op_tests/flydsl_tests/ragged_nn.py | 103 +-- .../test_flydsl_fp8_paged_mqa_logits.py | 657 +----------------- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 2 +- ...t_flydsl_fp8_paged_mqa_logits_ragged_nn.py | 4 +- op_tests/test_flydsl_fp8_paged_mqa_logits.py | 632 +++++++++++++++++ 5 files changed, 664 insertions(+), 734 deletions(-) create mode 100644 op_tests/test_flydsl_fp8_paged_mqa_logits.py diff --git a/op_tests/flydsl_tests/ragged_nn.py b/op_tests/flydsl_tests/ragged_nn.py index c18be251894..b35ce1f0d93 100644 --- a/op_tests/flydsl_tests/ragged_nn.py +++ b/op_tests/flydsl_tests/ragged_nn.py @@ -1,98 +1,13 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Compact or ragged-``next_n`` batch for paged FP8 MQA logits. - -Q is dense ``[B, next_n, H, D]``. For a uniform batch all rows are live, like -Gluon. For a ragged batch, ``next_n_lens[b]`` in ``{1..next_n}`` marks how many -leading rows are live. The gfx950 FlyDSL kernel consumes that optional vector; -unused Q/weight rows are ignored and unused logit rows stay ``-inf``. -""" - -from __future__ import annotations - -import torch - -from aiter.ops.flydsl.kernels.mqa_logits.fp8_paged_mqa_logits_gfx950 import MAX_NN -from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( - ref_fp8_paged_mqa_logits, +"""Compact or ragged-``next_n`` helpers. Implementation lives on the op test.""" + +from op_tests.test_flydsl_fp8_paged_mqa_logits import ( # noqa: F401 + MAX_NN, + live_row_mask, + padded_logits, + ref_padded_ragged, + sample_next_n_lens, + scatter_seq_logits, ) - - -def sample_next_n_lens(nq: int, max_nn: int = MAX_NN, seed: int = 1079) -> torch.Tensor: - g = torch.Generator(device="cpu") - g.manual_seed(seed + nq) - return torch.randint( - 1, - max_nn + 1, - (nq,), - generator=g, - dtype=torch.int32, - device="cpu", - ) - - -def padded_logits(batch: int, max_nn: int, max_model_len: int) -> torch.Tensor: - return torch.full( - (batch * max_nn, max_model_len), - float("-inf"), - dtype=torch.float32, - device="cuda", - ) - - -def scatter_seq_logits( - out: torch.Tensor, - seq: int, - max_nn: int, - row_logits: torch.Tensor, - n_rows: int, -) -> None: - """Copy ``n_rows`` of a compact ``[n, max_model_len]`` tensor into padded out.""" - dst = seq * max_nn - out[dst : dst + n_rows].copy_(row_logits[:n_rows]) - - -def ref_padded_ragged( - q_pad: torch.Tensor, - kv_cache_fp8: torch.Tensor, - weights_pad: torch.Tensor, - context_lens: torch.Tensor, - block_tables: torch.Tensor, - next_n_lens: torch.Tensor, - max_model_len: int, - fp8_dtype, - *, - max_nn: int = MAX_NN, - block_size: int = 64, -) -> torch.Tensor: - """Per-sequence compact ref, scattered into ``[B*max_nn, max_model_len]``.""" - batch = q_pad.shape[0] - out = padded_logits(batch, max_nn, max_model_len) - nn = next_n_lens.tolist() - for b, n in enumerate(nn): - n = int(n) - q_b = q_pad[b : b + 1, :n] - w_b = weights_pad[b * max_nn : b * max_nn + n] - compact = ref_fp8_paged_mqa_logits( - q_b, - kv_cache_fp8, - w_b, - context_lens[b : b + 1], - block_tables[b : b + 1], - max_model_len, - fp8_dtype, - block_size=block_size, - ) - scatter_seq_logits(out, b, max_nn, compact, n) - return out - - -def live_row_mask( - next_n_lens: torch.Tensor, max_nn: int, max_model_len: int -) -> torch.Tensor: - """True on live ``[B*max_nn, max_model_len]`` rows (all columns).""" - batch = next_n_lens.shape[0] - rows = torch.arange(max_nn, device=next_n_lens.device) - live = rows[None, :] < next_n_lens[:, None] - return live.reshape(batch * max_nn, 1).expand(-1, max_model_len) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py index e492b30abc3..5647e710db8 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -1,643 +1,28 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness sweep and benchmarks for the FlyDSL paged FP8 MQA-logits kernel. - -The reference is a torch port of vLLM's ``fp8_paged_mqa_logits_torch``, copied -here so the test does not depend on vLLM. Gate: exact ``-inf``-mask match plus -``calc_diff < 1e-3``, tolerances NOT widened. - -Performance timing and optional Gluon A/B comparison live in this file rather -than separate benchmark scripts. -""" - -import argparse -import itertools -import math -import random -from typing import NamedTuple - -import pandas as pd -import torch - -import aiter -from aiter import dtypes -from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits -from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from aiter.test_common import benchmark, checkAllclose, run_perftest - -torch.set_default_device("cuda") - -SUPPORTED_GFX = ["gfx950"] -_E4M3_NATIVE = get_fp8_e4m3_dtype() -DTYPE_MAP = {"fnuz": _E4M3_NATIVE, "fn": torch.float8_e4m3fn} - -REF_CHUNK_K = 256 -REF_WAVE_PER_EU = 2 - - -def calc_diff(x, y): - x, y = x.double(), y.double() - denominator = (x * x + y * y).sum() - return 1 - 2 * (x * y).sum() / denominator - - -def kv_cache_cast_to_fp8(x, fp8_dtype): - """Co-pack a bf16 KV cache into the fp8+scale byte layout: per block-row, - KVBlockSize*head_dim fp8 bytes then KVBlockSize f32 scales, no padding. - """ - num_blocks, block_size, num_heads, head_dim = x.shape - assert num_heads == 1 - x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) - sf = x_amax / 240.0 - x_scaled = (x * (1.0 / sf)).to(fp8_dtype) - x_fp8 = torch.empty( - (num_blocks, block_size * (head_dim + 4)), - device=x.device, - dtype=torch.uint8, - ) - x_fp8[:, : block_size * head_dim] = x_scaled.view( - num_blocks, block_size * head_dim - ).view(dtype=torch.uint8) - x_fp8[:, block_size * head_dim : block_size * head_dim + 4 * block_size] = sf.view( - num_blocks, block_size - ).view(dtype=torch.uint8) - return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) - - -def preshuffle_kv_data(kv_cache_fp8, head_dim): - """Apply ``shuffle_weight(layout=(16,16))`` to the per-block fp8 key bytes, - leaving the co-packed f32 scale tail alone -- the production Preshuffle - layout. Only the kernel gets this copy; the oracle reads the unshuffled - cache, since the reference is layout-agnostic. - """ - from aiter.ops.shuffle import shuffle_weight - - num_blocks, block_size, one, index_dim = kv_cache_fp8.shape - assert block_size % 16 == 0, "preshuffle requires KVBlockSize % 16 == 0" - flat = kv_cache_fp8.reshape(num_blocks, block_size * index_dim).clone() - data = ( - flat[:, : block_size * head_dim] - .contiguous() - .view(num_blocks, block_size, head_dim) - ) - shuffled = shuffle_weight(data, layout=(16, 16)).reshape( - num_blocks, block_size * head_dim - ) - flat[:, : block_size * head_dim] = shuffled - return flat.view(num_blocks, block_size, one, index_dim) - - -def ref_fp8_paged_mqa_logits( - q, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - block_size=1, -): - """Torch reference, vectorized port of vLLM ``fp8_paged_mqa_logits_torch``.""" - batch_size, next_n, _heads, dim = q.size() - num_blocks = kv_cache_fp8.shape[0] - index_dim = kv_cache_fp8.shape[-1] - flat = kv_cache_fp8.reshape(num_blocks, block_size * index_dim) - keys = ( - flat[:, : block_size * dim] - .contiguous() - .view(fp8_dtype) - .float() - .view(num_blocks, block_size, dim) - ) - scales = ( - flat[:, block_size * dim : block_size * dim + 4 * block_size] - .contiguous() - .view(torch.float32) - .view(num_blocks, block_size, 1) - ) - kvf = keys * scales - qf = q.float() - logits = torch.full( - [batch_size * next_n, max_model_len], - float("-inf"), - device=q.device, - dtype=torch.float32, - ) - for i in range(batch_size): - context_len = int(context_lens[i].item()) - if context_len == 0: - continue - pos = torch.arange(context_len, device=q.device) - blk = block_tables[i, pos // block_size] - tok = pos % block_size - kx = kvf[blk, tok] - s = torch.einsum("nhd,pd->nhp", qf[i], kx) - s = torch.relu(s) - wl = weights[i * next_n : (i + 1) * next_n, :] - s = (s * wl[:, :, None]).sum(dim=1) - q_lim = ( - context_len - next_n + torch.arange(next_n, device=q.device) - ).unsqueeze(1) - s = torch.where(pos[None, :] <= q_lim, s, float("-inf")) - logits[i * next_n : (i + 1) * next_n, :context_len] = s - return logits - - -class Inputs(NamedTuple): - q: torch.Tensor - q_fp8: torch.Tensor - kv_cache_fp8: torch.Tensor - weights: torch.Tensor - context_lens: torch.Tensor - block_tables: torch.Tensor - max_model_len: int - fp8_dtype: torch.dtype - - -def _build_inputs( - batch_size, - next_n, - heads, - head_dim, - avg_kv_length, - q_dtype, - block_size=1, - seed=0, - var_ratio=0.0, - pool_blocks=0, -): - torch.manual_seed(seed) - random.seed(seed) - fp8_dtype = get_fp8_e4m3_dtype() - - max_model_len = 2 * avg_kv_length - if var_ratio == 0.0: - context_lens = torch.full( - (batch_size,), avg_kv_length, device="cuda", dtype=torch.int32 - ) - else: - lo = max(1, int((1 - var_ratio) * avg_kv_length)) - hi = int((1 + var_ratio) * avg_kv_length) + 1 - context_lens = torch.randint(lo, hi, (batch_size,)).cuda().to(torch.int32) - context_lens = torch.clamp(context_lens, min=next_n) - - blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size - max_block_len = int(blocks_per_seq.max().item()) - needed_blocks = int(blocks_per_seq.sum().item()) - num_blocks = needed_blocks if pool_blocks <= 0 else max(pool_blocks, max_block_len) - - q = torch.randn((batch_size, next_n, heads, head_dim), dtype=torch.bfloat16) - kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) - weights = torch.randn((batch_size * next_n, heads), dtype=torch.float32) - - pool = list(range(num_blocks)) - random.shuffle(pool) - pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) - col = torch.arange(max_block_len, device="cuda", dtype=torch.int64) - starts = torch.cumsum(blocks_per_seq, 0) - blocks_per_seq - block_tables = torch.where( - col[None, :] < blocks_per_seq[:, None], - pool_t[(starts[:, None] + col[None, :]) % num_blocks], - torch.zeros((), device="cuda", dtype=torch.int32), - ).to(torch.int32) - - q_fp8 = q.to(q_dtype) - kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, fp8_dtype) - return Inputs( - q, - q_fp8, - kv_cache_fp8, - weights, - context_lens, - block_tables, - max_model_len, - fp8_dtype, - ) - - -def _kernel_inputs(inp, batch_size, next_n, head_dim): - kv_cache_kernel = preshuffle_kv_data(inp.kv_cache_fp8, head_dim) - out = torch.full( - (batch_size * next_n, inp.max_model_len), - float("-inf"), - device="cuda", - dtype=torch.float32, - ) - return kv_cache_kernel, out - - -def _verify_paged_mqa_logits( - inp, - batch_size, - next_n, - head_dim, - split_kv, - block_size, - preshuffle, -): - _split_kv = None if split_kv == 0 else split_kv - kv_cache_kernel, out = _kernel_inputs(inp, batch_size, next_n, head_dim) - - with torch.inference_mode(): - ref = ref_fp8_paged_mqa_logits( - inp.q, - inp.kv_cache_fp8, - inp.weights, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - inp.fp8_dtype, - block_size=block_size, - ) - ref_mask = ref == float("-inf") - - with torch.inference_mode(): - got = flydsl_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache_kernel, - inp.weights, - out, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - Preshuffle=preshuffle, - KVBlockSize=block_size, - SplitKV=_split_kv, - ) - - got_mask = got == float("-inf") - assert torch.equal(got_mask, ref_mask), "flydsl paged: -inf mask mismatch" - - err = 0.0 - if not ref_mask.all(): - diff = calc_diff(got.masked_fill(got_mask, 0), ref.masked_fill(ref_mask, 0)) - assert diff < 1e-3, f"flydsl paged calc_diff={diff}" - err = diff.item() - checkAllclose( - ref.masked_fill(ref_mask, 0).to(dtypes.fp32), - got.masked_fill(got_mask, 0).to(dtypes.fp32), - rtol=1e-2, - atol=5.0, - msg="flydsl paged fp8_mqa_logits", - printLog=False, - ) - return err, kv_cache_kernel, out - - -@benchmark() -def _bench_flydsl_paged_kernel( - q_fp8, - kv_cache_kernel, - weights, - out, - context_lens, - block_tables, - max_model_len, - *, - preshuffle, - block_size, - split_kv, -): - _split_kv = None if split_kv == 0 else split_kv - - def fn(): - return flydsl_fp8_paged_mqa_logits( - q_fp8, - kv_cache_kernel, - weights, - out, - context_lens, - block_tables, - max_model_len, - Preshuffle=preshuffle, - KVBlockSize=block_size, - SplitKV=_split_kv, - ) - - with torch.inference_mode(): - _, us = run_perftest(fn) - return {"flydsl us": float(us)} - - -def run_case( - batch_size, - next_n, - heads, - head_dim, - avg_kv_length, - q_dtype, - split_kv=0, - block_size=64, - preshuffle=True, - bench=False, -): - inp = _build_inputs( - batch_size, - next_n, - heads, - head_dim, - avg_kv_length, - DTYPE_MAP[q_dtype], - block_size=block_size, - ) - err, kv_cache_kernel, out = _verify_paged_mqa_logits( - inp, - batch_size, - next_n, - head_dim, - split_kv, - block_size, - preshuffle, - ) - ret = { - "gfx": get_gfx(), - "kvb": block_size, - "preshuffle": preshuffle, - "split_kv": "auto" if split_kv == 0 else split_kv, - "flydsl err": err, - } - if not bench: - return ret - - ret.update( - _bench_flydsl_paged_kernel( - inp.q_fp8, - kv_cache_kernel, - inp.weights, - out, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - preshuffle=preshuffle, - block_size=block_size, - split_kv=split_kv, - ) - ) - - from aiter.ops.triton.attention.pa_mqa_logits import ( - deepgemm_fp8_paged_mqa_logits, - ) - - out_ref = torch.full( - (batch_size * next_n, inp.max_model_len), - float("-inf"), - device="cuda", - dtype=torch.float32, - ) - - def fn_ref(): - return deepgemm_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache_kernel, - inp.weights, - out_ref, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - ChunkK=REF_CHUNK_K, - Preshuffle=True, - KVBlockSize=block_size, - WavePerEU=REF_WAVE_PER_EU, - ) - - with torch.inference_mode(): - _, ref_us = run_perftest(fn_ref) - ret["gluon us"] = float(ref_us) - - return ret - - -_BASE = { - "batch_size": 2, - "next_n": 8, - "heads": 32, - "head_dim": 128, - "avg_kv_length": 1024, - "q_dtype": "fnuz", -} -_AXES = ( - "batch_size", - "next_n", - "heads", - "head_dim", - "avg_kv_length", - "q_dtype", - "split_kv", - "block_size", +"""Shim for WaveScope drivers. Canonical test: op_tests/test_flydsl_fp8_paged_mqa_logits.py.""" + +from op_tests.test_flydsl_fp8_paged_mqa_logits import ( # noqa: F401 + HEAD_DIM, + HEADS, + KV_BLOCK_SIZE, + MAX_NN, + Inputs, + _build_inputs, + _kernel_inputs, + calc_diff, + kv_cache_cast_to_fp8, + live_row_mask, + main, + padded_logits, + preshuffle_kv_data, + ref_fp8_paged_mqa_logits, + ref_padded_ragged, + run_torch, + sample_next_n_lens, + scatter_seq_logits, ) - -def _c(**kw): - return {**_BASE, "block_size": 64, "preshuffle": True, **kw} - - -def default_cases(): - cases = [_c(batch_size=b) for b in (1, 2, 4, 8)] - cases += [_c(avg_kv_length=kv) for kv in (128, 8192)] - cases += [_c(split_kv=sk) for sk in (1, 4)] - cases.append(_c(avg_kv_length=128, split_kv=4)) - cases.append(_c(avg_kv_length=8192, split_kv=1)) - return cases - - -def exhaustive_cases(): - prod = itertools.product( - [(1, 8), (2, 8), (4, 8), (8, 8)], - [32], - [128], - [128, 1024, 8192], - ["fnuz"], - [0, 1, 4], - [64], - ) - return [ - {**dict(zip(_AXES, (bs, nn, nh, hd, kv, qd, sk, kvb))), "preshuffle": True} - for (bs, nn), nh, hd, kv, qd, sk, kvb in prod - ] - - -def gluon_ab_shapes(): - H, D = 32, 128 - shapes = [] - for B in (1, 4, 16, 64, 128): - for avg_kv in (16384, 32768, 65536): - shapes.append((B, 2, H, D, avg_kv)) - for B in (1, 16, 128): - shapes.append((B, 1, H, D, 32768)) - return shapes - - -def run_profile(args): - inp = _build_inputs( - args.batch, - args.next_n, - args.heads, - args.head_dim, - args.kv_len, - DTYPE_MAP[args.q_dtype], - block_size=args.kv_block_size, - seed=args.seed, - var_ratio=args.var_ratio, - pool_blocks=args.pool_blocks, - ) - err, kv_cache_kernel, out = _verify_paged_mqa_logits( - inp, - args.batch, - args.next_n, - args.head_dim, - args.split_kv, - args.kv_block_size, - True, - ) - _split_kv = None if args.split_kv == 0 else args.split_kv - - def run(): - flydsl_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache_kernel, - inp.weights, - out, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - Preshuffle=True, - KVBlockSize=args.kv_block_size, - SplitKV=_split_kv, - ) - - run() - torch.cuda.synchronize() - total_ctx = int(inp.context_lens.sum().item()) - kv_bytes = args.next_n * total_ctx * (args.head_dim + 4) - print(f"# profile correctness err={err:.3e}") - print( - f"# gfx={get_gfx()} B={args.batch} nn={args.next_n} H={args.heads} " - f"D={args.head_dim} kv_len={args.kv_len} kvb={args.kv_block_size} " - f"preshuffle=True iters={args.iters}" - ) - print(f"# KV requested: {kv_bytes / 1e6:.1f} MB") - torch.cuda.synchronize() - for _ in range(args.iters): - run() - torch.cuda.synchronize() - if args.time: - _, us = run_perftest(run, num_iters=args.num_iters) - flops = 2 * args.heads * args.head_dim * args.next_n * total_ctx - print( - f"time: {us:.2f} us | {flops / us / 1e6:.1f} TFLOP/s | " - f"{kv_bytes / us / 1e3:.1f} GB/s requested" - ) - - -def run_gluon_ab(args): - rows = [] - for B, nn, H, D, kv_len in gluon_ab_shapes(): - ret = run_case( - batch_size=B, - next_n=nn, - heads=H, - head_dim=D, - avg_kv_length=kv_len, - q_dtype="fnuz", - block_size=64, - preshuffle=True, - bench=True, - ) - fly_us = ret["flydsl us"] - ref_us = ret.get("gluon us", float("nan")) - rows.append( - { - **ret, - "B": B, - "nn": nn, - "avg_kv": kv_len, - "fly/ref": ( - fly_us / ref_us - if not math.isnan(ref_us) and ref_us > 0 - else float("nan") - ), - } - ) - print( - f"B={B} nn={nn} kv={kv_len} fly={fly_us:.1f}us gluon={ref_us:.1f}us " - f"fly/ref={rows[-1]['fly/ref']:.2f}x err={ret['flydsl err']:.2e}", - flush=True, - ) - df = pd.DataFrame(rows) - try: - summary = df.to_markdown(index=False) - except ImportError: - summary = df.to_string(index=False) - aiter.logger.info("gluon A/B summary:\n%s", summary) - - -def main(): - if get_gfx() not in SUPPORTED_GFX: - aiter.logger.warning( - "fp8_paged_mqa_logits unsupported on %s; skipping", get_gfx() - ) - return - - parser = argparse.ArgumentParser( - description="FlyDSL paged fp8_mqa_logits correctness + benchmarks" - ) - parser.add_argument( - "--exhaustive", - action="store_true", - help="run the full cartesian product instead of the curated matrix", - ) - parser.add_argument( - "--bench", - action="store_true", - help="also run run_perftest timing in the correctness sweep", - ) - parser.add_argument( - "--compare-gluon", - action="store_true", - help="run the Gluon A/B decode shape sweep (KVB=64, preshuffle)", - ) - parser.add_argument( - "--profile", - action="store_true", - help="run a single-config timed loop (for rocprofv3 / isolation profiling)", - ) - parser.add_argument("--batch", type=int, default=16) - parser.add_argument("--next-n", type=int, default=8) - parser.add_argument("--heads", type=int, default=32) - parser.add_argument("--head-dim", type=int, default=128) - parser.add_argument("--kv-len", type=int, default=32768) - parser.add_argument("--kv-block-size", type=int, default=64) - parser.add_argument("--q-dtype", type=str, default="fnuz", choices=["fnuz", "fn"]) - parser.add_argument("--split-kv", type=int, default=0) - parser.add_argument("--var-ratio", type=float, default=0.0) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--pool-blocks", type=int, default=0) - parser.add_argument("--iters", type=int, default=100) - parser.add_argument("--num-iters", type=int, default=101) - parser.add_argument("--time", action="store_true") - args = parser.parse_args() - - if args.profile: - if args.kv_block_size != 64 or args.heads != 32 or args.next_n != 8: - raise SystemExit("kernel requires H=32, padded next_n=8, KVB=64") - run_profile(args) - return - if args.compare_gluon: - run_gluon_ab(args) - return - - cases = exhaustive_cases() if args.exhaustive else default_cases() - - df = [run_case(bench=args.bench, **c) for c in cases] - df = pd.DataFrame(df) - try: - summary = df.to_markdown(index=False) - except ImportError: - summary = df.to_string(index=False) - aiter.logger.info("fp8_paged_mqa_logits summary:\n%s", summary) - - if __name__ == "__main__": main() diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index b1d99e4b83a..b837c794d0f 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -17,7 +17,7 @@ ref_padded_ragged, sample_next_n_lens, ) -from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( +from op_tests.test_flydsl_fp8_paged_mqa_logits import ( _build_inputs, _kernel_inputs, calc_diff, diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py index bb9e7ce8bdb..d42f7c60cd9 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py @@ -7,9 +7,7 @@ import torch from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl.kernels.mqa_logits.fp8_paged_mqa_logits_gfx950 import ( - flydsl_fp8_paged_mqa_logits, -) +from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype from op_tests.flydsl_tests.ragged_nn import ( MAX_NN, diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py new file mode 100644 index 00000000000..48ba83634d9 --- /dev/null +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -0,0 +1,632 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""FlyDSL paged FP8 MQA logits (decode) vs a torch reference. + +The kernel is the gfx950 H=32 / D=128 / KVB=64 preshuffled path the indexer +calls: compact Q ``[B, next_n, 32, 128]``, preallocated ``out_logits``, optional +``next_n_lens``. Torch is the reference only (not timed). Gluon is a candidate +on the uniform compact path. +""" + +from __future__ import annotations + +import argparse +import itertools +import random +from typing import NamedTuple + +import pandas as pd +import torch + +import aiter +from aiter import dtypes +from aiter.jit.utils.chip_info import get_gfx +from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits +from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype +from aiter.test_common import benchmark, checkAllclose, run_perftest + +torch.set_default_device("cuda") + +SUPPORTED_GFX = ["gfx950"] +HEADS = 32 +HEAD_DIM = 128 +KV_BLOCK_SIZE = 64 +INDEX_DIM = HEAD_DIM + 4 +MAX_NN = 8 +WIDE_MAX_MODEL_LEN = 1 << 20 +_E4M3_NATIVE = get_fp8_e4m3_dtype() + +try: + from aiter.ops.triton.attention.pa_mqa_logits import ( + deepgemm_fp8_paged_mqa_logits, + ) +except ImportError: + deepgemm_fp8_paged_mqa_logits = None + + +def calc_diff(x, y): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + return 1 - 2 * (x * y).sum() / denominator + + +def kv_cache_cast_to_fp8(x, fp8_dtype): + num_blocks, block_size, num_heads, head_dim = x.shape + assert num_heads == 1 + x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) + sf = x_amax / 240.0 + x_scaled = (x * (1.0 / sf)).to(fp8_dtype) + x_fp8 = torch.empty( + (num_blocks, block_size * (head_dim + 4)), + device=x.device, + dtype=torch.uint8, + ) + x_fp8[:, : block_size * head_dim] = x_scaled.view( + num_blocks, block_size * head_dim + ).view(dtype=torch.uint8) + x_fp8[:, block_size * head_dim : block_size * head_dim + 4 * block_size] = sf.view( + num_blocks, block_size + ).view(dtype=torch.uint8) + return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4) + + +def preshuffle_kv_data(kv_cache_fp8, head_dim): + from aiter.ops.shuffle import shuffle_weight + + num_blocks, block_size, one, index_dim = kv_cache_fp8.shape + assert block_size % 16 == 0, "preshuffle requires KVBlockSize % 16 == 0" + flat = kv_cache_fp8.reshape(num_blocks, block_size * index_dim).clone() + data = ( + flat[:, : block_size * head_dim] + .contiguous() + .view(num_blocks, block_size, head_dim) + ) + shuffled = shuffle_weight(data, layout=(16, 16)).reshape( + num_blocks, block_size * head_dim + ) + flat[:, : block_size * head_dim] = shuffled + return flat.view(num_blocks, block_size, one, index_dim) + + +def run_torch( + q, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + block_size=KV_BLOCK_SIZE, +): + """Torch reference. Not timed, not in the table.""" + batch_size, next_n, _heads, dim = q.size() + num_blocks = kv_cache_fp8.shape[0] + index_dim = kv_cache_fp8.shape[-1] + flat = kv_cache_fp8.reshape(num_blocks, block_size * index_dim) + keys = ( + flat[:, : block_size * dim] + .contiguous() + .view(fp8_dtype) + .float() + .view(num_blocks, block_size, dim) + ) + scales = ( + flat[:, block_size * dim : block_size * dim + 4 * block_size] + .contiguous() + .view(torch.float32) + .view(num_blocks, block_size, 1) + ) + kvf = keys * scales + qf = q.float() + logits = torch.full( + [batch_size * next_n, max_model_len], + float("-inf"), + device=q.device, + dtype=torch.float32, + ) + for i in range(batch_size): + context_len = int(context_lens[i].item()) + if context_len == 0: + continue + pos = torch.arange(context_len, device=q.device) + blk = block_tables[i, pos // block_size] + tok = pos % block_size + kx = kvf[blk, tok] + s = torch.einsum("nhd,pd->nhp", qf[i], kx) + s = torch.relu(s) + wl = weights[i * next_n : (i + 1) * next_n, :] + s = (s * wl[:, :, None]).sum(dim=1) + q_lim = ( + context_len - next_n + torch.arange(next_n, device=q.device) + ).unsqueeze(1) + s = torch.where(pos[None, :] <= q_lim, s, float("-inf")) + logits[i * next_n : (i + 1) * next_n, :context_len] = s + return logits + + +ref_fp8_paged_mqa_logits = run_torch + + +class Inputs(NamedTuple): + q: torch.Tensor + q_fp8: torch.Tensor + kv_cache_fp8: torch.Tensor + weights: torch.Tensor + context_lens: torch.Tensor + block_tables: torch.Tensor + max_model_len: int + fp8_dtype: torch.dtype + + +def _build_inputs( + batch_size, + next_n, + heads, + head_dim, + avg_kv_length, + q_dtype, + block_size=KV_BLOCK_SIZE, + seed=0, + var_ratio=0.0, + pool_blocks=0, +): + torch.manual_seed(seed) + random.seed(seed) + fp8_dtype = get_fp8_e4m3_dtype() + + max_model_len = 2 * avg_kv_length + if var_ratio == 0.0: + context_lens = torch.full( + (batch_size,), avg_kv_length, device="cuda", dtype=torch.int32 + ) + else: + lo = max(1, int((1 - var_ratio) * avg_kv_length)) + hi = int((1 + var_ratio) * avg_kv_length) + 1 + context_lens = torch.randint(lo, hi, (batch_size,)).cuda().to(torch.int32) + context_lens = torch.clamp(context_lens, min=next_n) + + blocks_per_seq = (context_lens.to(torch.int64) + block_size - 1) // block_size + max_block_len = int(blocks_per_seq.max().item()) + needed_blocks = int(blocks_per_seq.sum().item()) + num_blocks = needed_blocks if pool_blocks <= 0 else max(pool_blocks, max_block_len) + + q = torch.randn((batch_size, next_n, heads, head_dim), dtype=torch.bfloat16) + kv_cache = torch.randn((num_blocks, block_size, 1, head_dim), dtype=torch.bfloat16) + weights = torch.randn((batch_size * next_n, heads), dtype=torch.float32) + + pool = list(range(num_blocks)) + random.shuffle(pool) + pool_t = torch.tensor(pool, device="cuda", dtype=torch.int32) + col = torch.arange(max_block_len, device="cuda", dtype=torch.int64) + starts = torch.cumsum(blocks_per_seq, 0) - blocks_per_seq + block_tables = torch.where( + col[None, :] < blocks_per_seq[:, None], + pool_t[(starts[:, None] + col[None, :]) % num_blocks], + torch.zeros((), device="cuda", dtype=torch.int32), + ).to(torch.int32) + + q_fp8 = q.to(q_dtype) + kv_cache_fp8 = kv_cache_cast_to_fp8(kv_cache, fp8_dtype) + return Inputs( + q, + q_fp8, + kv_cache_fp8, + weights, + context_lens, + block_tables, + max_model_len, + fp8_dtype, + ) + + +def _kernel_inputs(inp, batch_size, next_n, head_dim): + kv_cache_kernel = preshuffle_kv_data(inp.kv_cache_fp8, head_dim) + out = torch.full( + (batch_size * next_n, inp.max_model_len), + float("-inf"), + device="cuda", + dtype=torch.float32, + ) + return kv_cache_kernel, out + + +def sample_next_n_lens(nq: int, max_nn: int = MAX_NN, seed: int = 1079) -> torch.Tensor: + g = torch.Generator(device="cpu") + g.manual_seed(seed + nq) + return torch.randint( + 1, + max_nn + 1, + (nq,), + generator=g, + dtype=torch.int32, + device="cpu", + ) + + +def padded_logits(batch: int, max_nn: int, max_model_len: int) -> torch.Tensor: + return torch.full( + (batch * max_nn, max_model_len), + float("-inf"), + dtype=torch.float32, + device="cuda", + ) + + +def scatter_seq_logits(out, seq, max_nn, row_logits, n_rows): + dst = seq * max_nn + out[dst : dst + n_rows].copy_(row_logits[:n_rows]) + + +def ref_padded_ragged( + q_pad, + kv_cache_fp8, + weights_pad, + context_lens, + block_tables, + next_n_lens, + max_model_len, + fp8_dtype, + *, + max_nn=MAX_NN, + block_size=KV_BLOCK_SIZE, +): + batch = q_pad.shape[0] + out = padded_logits(batch, max_nn, max_model_len) + for b, n in enumerate(next_n_lens.tolist()): + n = int(n) + compact = run_torch( + q_pad[b : b + 1, :n], + kv_cache_fp8, + weights_pad[b * max_nn : b * max_nn + n], + context_lens[b : b + 1], + block_tables[b : b + 1], + max_model_len, + fp8_dtype, + block_size=block_size, + ) + scatter_seq_logits(out, b, max_nn, compact, n) + return out + + +def live_row_mask(next_n_lens, max_nn, max_model_len): + batch = next_n_lens.shape[0] + rows = torch.arange(max_nn, device=next_n_lens.device) + live = rows[None, :] < next_n_lens[:, None] + return live.reshape(batch * max_nn, 1).expand(-1, max_model_len) + + +def _split_arg(split_kv): + return None if int(split_kv) <= 0 else int(split_kv) + + +def _grade(ref, got, msg): + ref_mask = ref == float("-inf") + got_mask = got == float("-inf") + assert torch.equal(got_mask, ref_mask), f"{msg}: -inf mask mismatch" + live = ~ref_mask + if not live.any(): + return 0.0 + diff = calc_diff(got.masked_fill(got_mask, 0), ref.masked_fill(ref_mask, 0)) + assert diff < 1e-3, f"{msg} calc_diff={diff}" + return checkAllclose( + ref.masked_fill(ref_mask, 0).to(dtypes.fp32), + got.masked_fill(got_mask, 0).to(dtypes.fp32), + rtol=1e-2, + atol=5.0, + msg=msg, + printLog=False, + ) + + +def _roofline(batch, next_n, kv_len, context_lens, q_fp8, weights, out): + total_ctx = int(context_lens.sum().item()) + flops = 2 * HEADS * HEAD_DIM * next_n * total_ctx + nbytes = ( + batch * next_n * HEADS * HEAD_DIM * q_fp8.element_size() + + total_ctx * INDEX_DIM + + batch * next_n * HEADS * weights.element_size() + + batch * next_n * kv_len * out.element_size() + ) + return flops, nbytes + + +def _time_candidates(candidates, ref, flops, nbytes, msg): + ret = {"gfx": get_gfx()} + for name, (fn, buf) in candidates.items(): + with torch.inference_mode(): + _, us = run_perftest(fn) + err = _grade(ref, buf, f"{name}: {msg}") + ret[f"{name} us"] = us + ret[f"{name} TFLOPS"] = flops / us / 1e6 + ret[f"{name} TB/s"] = nbytes / us / 1e6 + ret[f"{name} err"] = err + return ret + + +@benchmark() +def test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv): + inp = _build_inputs( + batch, next_n, HEADS, HEAD_DIM, kv_len, _E4M3_NATIVE, block_size=KV_BLOCK_SIZE + ) + kv_cache_kernel, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM) + ref = run_torch( + inp.q, + inp.kv_cache_fp8, + inp.weights, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + inp.fp8_dtype, + block_size=KV_BLOCK_SIZE, + ) + split = _split_arg(split_kv) + + def flydsl(): + out.fill_(float("-inf")) + return flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + SplitKV=split, + ) + + candidates = {"flydsl": (flydsl, out)} + # Gluon is a real kernel under test, not the torch reference. It only + # accepts a launch-wide next_n = Q.shape[1], so it is compact-only. + if deepgemm_fp8_paged_mqa_logits is not None: + out_g = out.clone() + + def gluon(): + out_g.fill_(float("-inf")) + return deepgemm_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out_g, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + ChunkK=256, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + WavePerEU=2, + ) + + candidates["gluon"] = (gluon, out_g) + + flops, nbytes = _roofline( + batch, next_n, kv_len, inp.context_lens, inp.q_fp8, inp.weights, out + ) + return _time_candidates(candidates, ref, flops, nbytes, "paged fp8_mqa_logits") + + +@benchmark() +def test_fp8_paged_mqa_logits_ragged(batch, next_n, kv_len, split_kv): + inp = _build_inputs( + batch, next_n, HEADS, HEAD_DIM, kv_len, _E4M3_NATIVE, block_size=KV_BLOCK_SIZE + ) + kv_cache_kernel, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM) + next_n_lens = sample_next_n_lens(batch, next_n).cuda() + ref = ref_padded_ragged( + inp.q, + inp.kv_cache_fp8, + inp.weights, + inp.context_lens, + inp.block_tables, + next_n_lens, + inp.max_model_len, + inp.fp8_dtype, + max_nn=next_n, + block_size=KV_BLOCK_SIZE, + ) + split = _split_arg(split_kv) + live = int(next_n_lens.sum().item()) + + def flydsl(): + out.fill_(float("-inf")) + return flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + next_n_lens=next_n_lens, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + SplitKV=split, + ) + + # Gluon omitted: it has no next_n_lens and would score padded rows. + flops, nbytes = _roofline( + batch, + live / max(batch, 1), + kv_len, + inp.context_lens, + inp.q_fp8, + inp.weights, + out, + ) + ret = _time_candidates( + {"flydsl": (flydsl, out)}, ref, flops, nbytes, "paged ragged" + ) + unused = ~live_row_mask(next_n_lens, next_n, inp.max_model_len)[:, 0] + if unused.any(): + assert torch.all(ref[unused] == float("-inf")) + return ret + + +@benchmark() +def test_fp8_paged_mqa_logits_wide_out(batch, kv_len): + """Padded output row 512 crosses the signed i32 byte-offset boundary.""" + next_n = MAX_NN + inp = _build_inputs( + batch, next_n, HEADS, HEAD_DIM, kv_len, _E4M3_NATIVE, block_size=KV_BLOCK_SIZE + ) + kv_cache_kernel, compact = _kernel_inputs(inp, batch, next_n, HEAD_DIM) + next_n_lens = torch.ones(batch, dtype=torch.int32, device="cuda") + wide = torch.full( + (batch * next_n, WIDE_MAX_MODEL_LEN), + float("-inf"), + dtype=torch.float32, + ) + split = 1 + + def flydsl(): + compact.fill_(float("-inf")) + wide.fill_(float("-inf")) + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + compact, + inp.context_lens, + inp.block_tables, + inp.max_model_len, + next_n_lens=next_n_lens, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + SplitKV=split, + ) + return flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + wide, + inp.context_lens, + inp.block_tables, + WIDE_MAX_MODEL_LEN, + next_n_lens=next_n_lens, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + SplitKV=split, + ) + + with torch.inference_mode(): + got, us = run_perftest(flydsl) + valid_cols = min(kv_len, compact.shape[1]) + assert torch.equal(got[:, :valid_cols], compact[:, :valid_cols]) + if batch * next_n > 512: + assert torch.equal(got[512:, :valid_cols], compact[512:, :valid_cols]) + flops, nbytes = _roofline( + batch, 1, kv_len, inp.context_lens, inp.q_fp8, inp.weights, compact + ) + return { + "gfx": get_gfx(), + "flydsl us": us, + "flydsl TFLOPS": flops / us / 1e6, + "flydsl TB/s": nbytes / us / 1e6, + "flydsl err": 0.0, + } + + +def _summarize(name, rows): + df = pd.DataFrame(rows) + try: + table = df.to_markdown(index=False) + except ImportError: + table = df.to_string(index=False) + aiter.logger.info("%s summary (markdown):\n%s", name, table) + + +def main(): + if get_gfx() not in SUPPORTED_GFX: + aiter.logger.warning( + "fp8_paged_mqa_logits unsupported on %s; skipping", get_gfx() + ) + return + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description="config input of test", + ) + parser.add_argument( + "-d", + "--dtype", + type=str, + nargs="*", + default=["fn"], + help="Q dtype tag. gfx950 uses native E4M3 FN.\n e.g.: -d fn", + ) + parser.add_argument( + "-b", + "--batch", + type=int, + nargs="*", + default=[1, 2, 4, 8], + help="Batch (sequences).\n e.g.: -b 4", + ) + parser.add_argument( + "--next-n", + type=int, + nargs="*", + default=[8], + help="Compact Q dim-1 / ragged pad.\n e.g.: --next-n 1 8", + ) + parser.add_argument( + "--kv-len", + type=int, + nargs="*", + default=[1024], + help="Tokens per sequence (multiple of kvb=64 preferred).\n e.g.: --kv-len 128 8192", + ) + parser.add_argument( + "--split-kv", + type=int, + nargs="*", + default=[0], + help="SplitKV; 0 = auto occupancy mapping.\n e.g.: --split-kv 0 1 4", + ) + parser.add_argument( + "--ragged-batch", + type=int, + nargs="*", + default=[4, 8], + help="Batch for the ragged table.", + ) + parser.add_argument( + "--ragged-kv-len", + type=int, + nargs="*", + default=[128, 192], + help="Context for the ragged table (short-page coverage).", + ) + parser.add_argument( + "--wide-batch", + type=int, + nargs="*", + default=[65], + help="Batch for the wide-output i32-offset table.", + ) + args = parser.parse_args() + + compact = [] + for _, batch, next_n, kv_len, split_kv in itertools.product( + args.dtype, args.batch, args.next_n, args.kv_len, args.split_kv + ): + compact.append(test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv)) + _summarize("fp8_paged_mqa_logits", compact) + + ragged = [] + for batch, next_n, kv_len, split_kv in itertools.product( + args.ragged_batch, args.next_n, args.ragged_kv_len, args.split_kv + ): + ragged.append(test_fp8_paged_mqa_logits_ragged(batch, next_n, kv_len, split_kv)) + _summarize("fp8_paged_mqa_logits_ragged", ragged) + + wide = [] + for batch in args.wide_batch: + wide.append(test_fp8_paged_mqa_logits_wide_out(batch, KV_BLOCK_SIZE)) + _summarize("fp8_paged_mqa_logits_wide_out", wide) + + +if __name__ == "__main__": + main() From 8098d25c585f691bda1eb7317963d249ec8c3b03 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Tue, 22 Sep 2026 19:13:41 +0000 Subject: [PATCH 51/52] test(flydsl): fold paged MQA pytest coverage into the op_test tables Keep a single gfx950-gated test_*.py as the source of truth, grade wide output against torch, and leave the flydsl_tests files as shims. Co-authored-by: Cursor --- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 168 +------------- ...t_flydsl_fp8_paged_mqa_logits_ragged_nn.py | 205 +----------------- op_tests/test_flydsl_fp8_paged_mqa_logits.py | 152 ++++++++----- 3 files changed, 107 insertions(+), 418 deletions(-) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py index b837c794d0f..5b97af07470 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py @@ -1,171 +1,9 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Focused long-context tests for the gfx950 compact/ragged LDS indexer.""" - -import argparse - -import pytest -import torch - -from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits -from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from aiter.test_common import run_perftest -from op_tests.flydsl_tests.ragged_nn import ( - MAX_NN, - ref_padded_ragged, - sample_next_n_lens, -) -from op_tests.test_flydsl_fp8_paged_mqa_logits import ( - _build_inputs, - _kernel_inputs, - calc_diff, -) - -torch.set_default_device("cuda") - -BATCH = 16 -HEADS = 32 -HEAD_DIM = 128 -KV_LEN = 32768 -KV_BLOCK_SIZE = 64 -WIDE_MAX_MODEL_LEN = 1 << 20 - - -def _inputs(kv_len=KV_LEN, batch=BATCH, seed=1079): - torch.manual_seed(seed) - inp = _build_inputs( - batch, - MAX_NN, - HEADS, - HEAD_DIM, - kv_len, - get_fp8_e4m3_dtype(), - block_size=KV_BLOCK_SIZE, - ) - kv_cache, out = _kernel_inputs(inp, batch, MAX_NN, HEAD_DIM) - next_n_lens = sample_next_n_lens(batch, MAX_NN, seed=seed).cuda() - return inp, kv_cache, out, next_n_lens - - -def _launch(inp, kv_cache, out, next_n_lens, *, max_model_len=None, split_kv=None): - return flydsl_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache, - inp.weights, - out, - inp.context_lens, - inp.block_tables, - inp.max_model_len if max_model_len is None else max_model_len, - next_n_lens=next_n_lens, - Preshuffle=True, - KVBlockSize=KV_BLOCK_SIZE, - SplitKV=split_kv, - ) - - -def _check(inp, kv_cache, out, next_n_lens, tag): - with torch.inference_mode(): - ref = ref_padded_ragged( - inp.q, - inp.kv_cache_fp8, - inp.weights, - inp.context_lens, - inp.block_tables, - next_n_lens, - inp.max_model_len, - inp.fp8_dtype, - max_nn=MAX_NN, - block_size=KV_BLOCK_SIZE, - ) - got = _launch(inp, kv_cache, out, next_n_lens) - - ref_mask = ref == float("-inf") - got_mask = got == float("-inf") - assert torch.equal(got_mask, ref_mask), f"{tag}: causal/padding -inf mask mismatch" - diff = calc_diff(got.masked_fill(got_mask, 0), ref.masked_fill(ref_mask, 0)) - assert diff < 1e-3, f"{tag} calc_diff={diff}" - - -@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") -def test_gfx950_ragged_long_context(): - inp, kv_cache, out, next_n_lens = _inputs() - _check(inp, kv_cache, out, next_n_lens, "ragged long context") - - -@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") -@pytest.mark.parametrize("pages", [2, 3, 4, 7]) -def test_gfx950_ragged_short_pages(pages): - inp, kv_cache, out, next_n_lens = _inputs( - kv_len=pages * KV_BLOCK_SIZE, batch=2, seed=pages - ) - _check(inp, kv_cache, out, next_n_lens, f"pages={pages}") - - -@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") -def test_gfx950_wide_output_no_tail_drop(): - """Padded output row 512 crosses the signed i32 byte-offset boundary.""" - batch = 65 - inp, kv_cache, compact, _ = _inputs(kv_len=KV_BLOCK_SIZE, batch=batch, seed=5) - next_n_lens = torch.ones(batch, dtype=torch.int32) - wide = torch.full( - (batch * MAX_NN, WIDE_MAX_MODEL_LEN), - float("-inf"), - dtype=torch.float32, - ) - - with torch.inference_mode(): - _launch(inp, kv_cache, compact, next_n_lens, split_kv=1) - _launch( - inp, - kv_cache, - wide, - next_n_lens, - max_model_len=WIDE_MAX_MODEL_LEN, - split_kv=1, - ) - torch.cuda.synchronize() - - valid_cols = KV_BLOCK_SIZE - assert torch.equal(wide[:, :valid_cols], compact[:, :valid_cols]) - assert torch.equal(wide[512:, :valid_cols], compact[512:, :valid_cols]) - - -def _benchmark(): - inp, kv_cache, out, next_n_lens = _inputs() - - def kernel(): - return _launch(inp, kv_cache, out, next_n_lens) - - with torch.inference_mode(): - _, us = run_perftest(kernel, num_iters=50, num_warmup=8) - print(f"time: {us:.3f} us") - - -def _profile(): - inp, kv_cache, out, next_n_lens = _inputs() - with torch.inference_mode(): - _launch(inp, kv_cache, out, next_n_lens) - torch.cuda.synchronize() - print("profile launch: pass") - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--profile", action="store_true") - parser.add_argument("--time", action="store_true") - args = parser.parse_args() - if args.time: - _benchmark() - elif args.profile: - _profile() - else: - test_gfx950_ragged_long_context() - for pages in (2, 3, 4, 7): - test_gfx950_ragged_short_pages(pages) - test_gfx950_wide_output_no_tail_drop() +"""WaveScope ATT driver. Canonical tables: op_tests/test_flydsl_fp8_paged_mqa_logits.py.""" +from op_tests import test_flydsl_fp8_paged_mqa_logits as _op if __name__ == "__main__": - main() + _op.test_fp8_paged_mqa_logits_ragged(16, _op.MAX_NN, 32768, 0) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py index d42f7c60cd9..5430e3ba763 100644 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py +++ b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py @@ -1,210 +1,15 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Correctness of compact and ragged next_n paged MQA logits.""" +"""Shim. Canonical test: op_tests/test_flydsl_fp8_paged_mqa_logits.py.""" -import pytest -import torch - -from aiter.jit.utils.chip_info import get_gfx -from aiter.ops.flydsl import flydsl_fp8_paged_mqa_logits -from aiter.ops.triton.utils.types import get_fp8_e4m3_dtype -from op_tests.flydsl_tests.ragged_nn import ( - MAX_NN, +from op_tests.test_flydsl_fp8_paged_mqa_logits import ( # noqa: F401 live_row_mask, + main, padded_logits, ref_padded_ragged, sample_next_n_lens, ) -from op_tests.flydsl_tests.test_flydsl_fp8_paged_mqa_logits import ( - calc_diff, - kv_cache_cast_to_fp8, - preshuffle_kv_data, -) - -torch.set_default_device("cuda") - -HEADS = 32 -HEAD_DIM = 128 -KVB = 64 - - -def _batch(nq, ctx, max_nn=MAX_NN, seed=0): - torch.manual_seed(seed) - fp8 = get_fp8_e4m3_dtype() - pages = ctx // KVB - q = torch.randn((nq, max_nn, HEADS, HEAD_DIM), dtype=torch.bfloat16) - weights = torch.randn((nq * max_nn, HEADS), dtype=torch.float32) - context_lens = torch.full((nq,), ctx, dtype=torch.int32) - kv = torch.randn((max(nq * pages, 1), KVB, 1, HEAD_DIM), dtype=torch.bfloat16) - kv_fp8 = kv_cache_cast_to_fp8(kv, fp8) - ids = torch.arange(nq * pages, device="cuda", dtype=torch.int32) - tables = torch.zeros((nq, pages), dtype=torch.int32) - tables[:, :pages] = ids.reshape(nq, pages) - return q, weights, context_lens, tables, kv_fp8, fp8 - - -def _run_and_check( - q, - weights, - context_lens, - tables, - kv_fp8, - fp8, - next_n_lens, - ctx, - *, - split_kv=None, -): - max_nn = q.shape[1] - ref = ref_padded_ragged( - q, - kv_fp8, - weights, - context_lens, - tables, - next_n_lens, - ctx, - fp8, - max_nn=max_nn, - block_size=KVB, - ) - out = padded_logits(q.shape[0], max_nn, ctx) - got = flydsl_fp8_paged_mqa_logits( - q.to(fp8), - preshuffle_kv_data(kv_fp8, HEAD_DIM), - weights, - out, - context_lens, - tables, - ctx, - next_n_lens=next_n_lens, - SplitKV=split_kv, - ) - ref_inf = ref == float("-inf") - got_inf = got == float("-inf") - assert torch.equal(got_inf, ref_inf), "ragged -inf mask mismatch" - diff = calc_diff(got.masked_fill(got_inf, 0), ref.masked_fill(ref_inf, 0)) - assert diff < 1e-3, diff - return got, ref - - -@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") -@pytest.mark.parametrize("nn", [1, 2, 4, 8]) -def test_compact_uniform_next_n(nn): - nq, ctx, max_nn = 4, 256, nn - next_n_lens = torch.full((nq,), nn, dtype=torch.int32, device="cuda") - q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, max_nn) - ref = ref_padded_ragged( - q, - kv_fp8, - weights, - context_lens, - tables, - next_n_lens, - ctx, - fp8, - max_nn=max_nn, - block_size=KVB, - ) - out = padded_logits(nq, max_nn, ctx) - got = flydsl_fp8_paged_mqa_logits( - q.to(fp8), - preshuffle_kv_data(kv_fp8, HEAD_DIM), - weights, - out, - context_lens, - tables, - ctx, - ) - ref_inf = ref == float("-inf") - got_inf = got == float("-inf") - assert torch.equal(got_inf, ref_inf) - diff = calc_diff(got.masked_fill(got_inf, 0), ref.masked_fill(ref_inf, 0)) - assert diff < 1e-3, diff - - -@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") -@pytest.mark.parametrize("max_nn", [4, MAX_NN]) -def test_ragged_mix_unused_stay_neg_inf(max_nn): - nq, ctx = 8, 192 - next_n_lens = sample_next_n_lens(nq, max_nn, seed=1079).cuda() - assert int(next_n_lens.min()) >= 1 and int(next_n_lens.max()) <= max_nn - assert len(next_n_lens.unique()) >= 2 - q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, max_nn, seed=1079) - got, ref = _run_and_check( - q, - weights, - context_lens, - tables, - kv_fp8, - fp8, - next_n_lens, - ctx, - ) - mask = live_row_mask(next_n_lens, max_nn, ctx) - unused_rows = ~mask[:, 0] - assert unused_rows.any(), "fixture should include padded rows" - assert torch.all(ref[unused_rows] == float("-inf")) - assert torch.all(got[unused_rows] == float("-inf")) - - q2 = torch.randn_like(q) - for b, n in enumerate(next_n_lens.tolist()): - q2[b, :n].copy_(q[b, :n]) - ref2 = ref_padded_ragged( - q2, - kv_fp8, - weights, - context_lens, - tables, - next_n_lens, - ctx, - fp8, - max_nn=max_nn, - block_size=KVB, - ) - assert torch.equal(ref == float("-inf"), ref2 == float("-inf")) - inf = ref == float("-inf") - diff = calc_diff(ref.masked_fill(inf, 0), ref2.masked_fill(inf, 0)) - assert diff < 1e-3, diff - - -@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") -@pytest.mark.parametrize("nq", [1, 24]) -def test_nq_and_short_pages(nq): - ctx, max_nn = 128, MAX_NN - next_n_lens = sample_next_n_lens(nq, max_nn, seed=3).cuda() - q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, max_nn, seed=3) - got, ref = _run_and_check( - q, - weights, - context_lens, - tables, - kv_fp8, - fp8, - next_n_lens, - ctx, - ) - assert ref.shape == (nq * max_nn, ctx) - assert got.shape == ref.shape - empty = padded_logits(nq, max_nn, ctx) - mask = live_row_mask(next_n_lens, max_nn, ctx) - assert torch.equal(ref[~mask[:, 0]], empty[~mask[:, 0]]) - -@pytest.mark.skipif(get_gfx() != "gfx950", reason="requires gfx950") -def test_odd_page_split_starts(): - nq, ctx = 3, 7 * KVB - next_n_lens = sample_next_n_lens(nq, MAX_NN, seed=19).cuda() - q, weights, context_lens, tables, kv_fp8, fp8 = _batch(nq, ctx, seed=19) - _run_and_check( - q, - weights, - context_lens, - tables, - kv_fp8, - fp8, - next_n_lens, - ctx, - split_kv=3, - ) +if __name__ == "__main__": + main() diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py index 48ba83634d9..9f905d58eec 100644 --- a/op_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -7,6 +7,11 @@ calls: compact Q ``[B, next_n, 32, 128]``, preallocated ``out_logits``, optional ``next_n_lens``. Torch is the reference only (not timed). Gluon is a candidate on the uniform compact path. + +Shape mapping (decode indexer): +- ``batch`` = nq sequences +- ``next_n`` = MTP rows per sequence (compact ``Q.shape[1]`` / ragged pad) +- ``kv_len`` = tokens per sequence (multiple of kvb=64 preferred) """ from __future__ import annotations @@ -36,6 +41,7 @@ MAX_NN = 8 WIDE_MAX_MODEL_LEN = 1 << 20 _E4M3_NATIVE = get_fp8_e4m3_dtype() +_Q_DTYPE = {"fn": torch.float8_e4m3fn} try: from aiter.ops.triton.attention.pa_mqa_logits import ( @@ -345,9 +351,10 @@ def _time_candidates(candidates, ref, flops, nbytes, msg): @benchmark() -def test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv): +def test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv, dtype): + q_dtype = _Q_DTYPE[dtype] inp = _build_inputs( - batch, next_n, HEADS, HEAD_DIM, kv_len, _E4M3_NATIVE, block_size=KV_BLOCK_SIZE + batch, next_n, HEADS, HEAD_DIM, kv_len, q_dtype, block_size=KV_BLOCK_SIZE ) kv_cache_kernel, out = _kernel_inputs(inp, batch, next_n, HEAD_DIM) ref = run_torch( @@ -461,6 +468,7 @@ def flydsl(): unused = ~live_row_mask(next_n_lens, next_n, inp.max_model_len)[:, 0] if unused.any(): assert torch.all(ref[unused] == float("-inf")) + assert torch.all(out[unused] == float("-inf")) return ret @@ -471,32 +479,32 @@ def test_fp8_paged_mqa_logits_wide_out(batch, kv_len): inp = _build_inputs( batch, next_n, HEADS, HEAD_DIM, kv_len, _E4M3_NATIVE, block_size=KV_BLOCK_SIZE ) - kv_cache_kernel, compact = _kernel_inputs(inp, batch, next_n, HEAD_DIM) + kv_cache_kernel, _compact = _kernel_inputs(inp, batch, next_n, HEAD_DIM) next_n_lens = torch.ones(batch, dtype=torch.int32, device="cuda") wide = torch.full( (batch * next_n, WIDE_MAX_MODEL_LEN), float("-inf"), dtype=torch.float32, ) + ref = ref_padded_ragged( + inp.q, + inp.kv_cache_fp8, + inp.weights, + inp.context_lens, + inp.block_tables, + next_n_lens, + inp.max_model_len, + inp.fp8_dtype, + max_nn=next_n, + block_size=KV_BLOCK_SIZE, + ) split = 1 + scored = wide[:, : inp.max_model_len] + def flydsl(): - compact.fill_(float("-inf")) wide.fill_(float("-inf")) flydsl_fp8_paged_mqa_logits( - inp.q_fp8, - kv_cache_kernel, - inp.weights, - compact, - inp.context_lens, - inp.block_tables, - inp.max_model_len, - next_n_lens=next_n_lens, - Preshuffle=True, - KVBlockSize=KV_BLOCK_SIZE, - SplitKV=split, - ) - return flydsl_fp8_paged_mqa_logits( inp.q_fp8, kv_cache_kernel, inp.weights, @@ -509,23 +517,14 @@ def flydsl(): KVBlockSize=KV_BLOCK_SIZE, SplitKV=split, ) + return scored - with torch.inference_mode(): - got, us = run_perftest(flydsl) - valid_cols = min(kv_len, compact.shape[1]) - assert torch.equal(got[:, :valid_cols], compact[:, :valid_cols]) - if batch * next_n > 512: - assert torch.equal(got[512:, :valid_cols], compact[512:, :valid_cols]) flops, nbytes = _roofline( - batch, 1, kv_len, inp.context_lens, inp.q_fp8, inp.weights, compact + batch, 1, kv_len, inp.context_lens, inp.q_fp8, inp.weights, wide + ) + return _time_candidates( + {"flydsl": (flydsl, scored)}, ref, flops, nbytes, "paged wide_out" ) - return { - "gfx": get_gfx(), - "flydsl us": us, - "flydsl TFLOPS": flops / us / 1e6, - "flydsl TB/s": nbytes / us / 1e6, - "flydsl err": 0.0, - } def _summarize(name, rows): @@ -554,6 +553,7 @@ def main(): type=str, nargs="*", default=["fn"], + choices=list(_Q_DTYPE), help="Q dtype tag. gfx950 uses native E4M3 FN.\n e.g.: -d fn", ) parser.add_argument( @@ -568,7 +568,7 @@ def main(): "--next-n", type=int, nargs="*", - default=[8], + default=[1, 2, 4, 8], help="Compact Q dim-1 / ragged pad.\n e.g.: --next-n 1 8", ) parser.add_argument( @@ -589,15 +589,43 @@ def main(): "--ragged-batch", type=int, nargs="*", - default=[4, 8], + default=[1, 8, 24], help="Batch for the ragged table.", ) + parser.add_argument( + "--ragged-next-n", + type=int, + nargs="*", + default=[8], + help="Ragged pad (max next_n).", + ) parser.add_argument( "--ragged-kv-len", type=int, nargs="*", - default=[128, 192], - help="Context for the ragged table (short-page coverage).", + default=[128, 192, 256, 448], + help="Context for the ragged table (2/3/4/7 pages of kvb=64).", + ) + parser.add_argument( + "--ragged-split-kv", + type=int, + nargs="*", + default=[0, 3], + help="SplitKV for the ragged table.", + ) + parser.add_argument( + "--long-batch", + type=int, + nargs="*", + default=[16], + help="Batch for the long-context ragged table.", + ) + parser.add_argument( + "--long-kv-len", + type=int, + nargs="*", + default=[32768], + help="Context for the long-context ragged table.", ) parser.add_argument( "--wide-batch", @@ -608,24 +636,42 @@ def main(): ) args = parser.parse_args() - compact = [] - for _, batch, next_n, kv_len, split_kv in itertools.product( - args.dtype, args.batch, args.next_n, args.kv_len, args.split_kv - ): - compact.append(test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv)) - _summarize("fp8_paged_mqa_logits", compact) - - ragged = [] - for batch, next_n, kv_len, split_kv in itertools.product( - args.ragged_batch, args.next_n, args.ragged_kv_len, args.split_kv - ): - ragged.append(test_fp8_paged_mqa_logits_ragged(batch, next_n, kv_len, split_kv)) - _summarize("fp8_paged_mqa_logits_ragged", ragged) - - wide = [] - for batch in args.wide_batch: - wide.append(test_fp8_paged_mqa_logits_wide_out(batch, KV_BLOCK_SIZE)) - _summarize("fp8_paged_mqa_logits_wide_out", wide) + compact = [ + test_fp8_paged_mqa_logits(batch, next_n, kv_len, split_kv, dtype) + for dtype, batch, next_n, kv_len, split_kv in itertools.product( + args.dtype, args.batch, args.next_n, args.kv_len, args.split_kv + ) + ] + if compact: + _summarize("fp8_paged_mqa_logits", compact) + + ragged = [ + test_fp8_paged_mqa_logits_ragged(batch, next_n, kv_len, split_kv) + for batch, next_n, kv_len, split_kv in itertools.product( + args.ragged_batch, + args.ragged_next_n, + args.ragged_kv_len, + args.ragged_split_kv, + ) + ] + if ragged: + _summarize("fp8_paged_mqa_logits_ragged", ragged) + + long_ctx = [ + test_fp8_paged_mqa_logits_ragged(batch, next_n, kv_len, split_kv) + for batch, next_n, kv_len, split_kv in itertools.product( + args.long_batch, args.ragged_next_n, args.long_kv_len, args.split_kv + ) + ] + if long_ctx: + _summarize("fp8_paged_mqa_logits_ragged_long", long_ctx) + + wide = [ + test_fp8_paged_mqa_logits_wide_out(batch, KV_BLOCK_SIZE) + for batch in args.wide_batch + ] + if wide: + _summarize("fp8_paged_mqa_logits_wide_out", wide) if __name__ == "__main__": From 5321501c48ebb1a0287b9d52344ecce0935e1374 Mon Sep 17 00:00:00 2001 From: Folkert Huizinga Date: Wed, 23 Sep 2026 08:11:10 +0000 Subject: [PATCH 52/52] test(flydsl): drop paged-MQA shims and add the probe_1079 decode table MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep a single op_test as the source of truth, with the production scattered-pool contract (mean Nq, next_n × ctx) graded against torch. Co-authored-by: Cursor --- op_tests/flydsl_tests/ragged_nn.py | 13 -- .../test_flydsl_fp8_paged_mqa_logits.py | 28 --- ...test_flydsl_fp8_paged_mqa_logits_gfx950.py | 9 - ...t_flydsl_fp8_paged_mqa_logits_ragged_nn.py | 15 -- op_tests/test_flydsl_fp8_paged_mqa_logits.py | 175 +++++++++++++++++- 5 files changed, 174 insertions(+), 66 deletions(-) delete mode 100644 op_tests/flydsl_tests/ragged_nn.py delete mode 100644 op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py delete mode 100644 op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py delete mode 100644 op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py diff --git a/op_tests/flydsl_tests/ragged_nn.py b/op_tests/flydsl_tests/ragged_nn.py deleted file mode 100644 index b35ce1f0d93..00000000000 --- a/op_tests/flydsl_tests/ragged_nn.py +++ /dev/null @@ -1,13 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Compact or ragged-``next_n`` helpers. Implementation lives on the op test.""" - -from op_tests.test_flydsl_fp8_paged_mqa_logits import ( # noqa: F401 - MAX_NN, - live_row_mask, - padded_logits, - ref_padded_ragged, - sample_next_n_lens, - scatter_seq_logits, -) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py deleted file mode 100644 index 5647e710db8..00000000000 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py +++ /dev/null @@ -1,28 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Shim for WaveScope drivers. Canonical test: op_tests/test_flydsl_fp8_paged_mqa_logits.py.""" - -from op_tests.test_flydsl_fp8_paged_mqa_logits import ( # noqa: F401 - HEAD_DIM, - HEADS, - KV_BLOCK_SIZE, - MAX_NN, - Inputs, - _build_inputs, - _kernel_inputs, - calc_diff, - kv_cache_cast_to_fp8, - live_row_mask, - main, - padded_logits, - preshuffle_kv_data, - ref_fp8_paged_mqa_logits, - ref_padded_ragged, - run_torch, - sample_next_n_lens, - scatter_seq_logits, -) - -if __name__ == "__main__": - main() diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py deleted file mode 100644 index 5b97af07470..00000000000 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py +++ /dev/null @@ -1,9 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""WaveScope ATT driver. Canonical tables: op_tests/test_flydsl_fp8_paged_mqa_logits.py.""" - -from op_tests import test_flydsl_fp8_paged_mqa_logits as _op - -if __name__ == "__main__": - _op.test_fp8_paged_mqa_logits_ragged(16, _op.MAX_NN, 32768, 0) diff --git a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py b/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py deleted file mode 100644 index 5430e3ba763..00000000000 --- a/op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py +++ /dev/null @@ -1,15 +0,0 @@ -# SPDX-License-Identifier: MIT -# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. - -"""Shim. Canonical test: op_tests/test_flydsl_fp8_paged_mqa_logits.py.""" - -from op_tests.test_flydsl_fp8_paged_mqa_logits import ( # noqa: F401 - live_row_mask, - main, - padded_logits, - ref_padded_ragged, - sample_next_n_lens, -) - -if __name__ == "__main__": - main() diff --git a/op_tests/test_flydsl_fp8_paged_mqa_logits.py b/op_tests/test_flydsl_fp8_paged_mqa_logits.py index 9f905d58eec..c638d8d6f76 100644 --- a/op_tests/test_flydsl_fp8_paged_mqa_logits.py +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -9,9 +9,14 @@ on the uniform compact path. Shape mapping (decode indexer): -- ``batch`` = nq sequences +- ``batch`` / ``nq`` = sequences - ``next_n`` = MTP rows per sequence (compact ``Q.shape[1]`` / ragged pad) - ``kv_len`` = tokens per sequence (multiple of kvb=64 preferred) + +``test_fp8_paged_mqa_logits_probe1079`` is the production decode contract from +the probe_1079 bs64 histogram: 44362-block KV pool, scattered 8 KiB pages, +``max_model_len=258048``, mean Nq=117. The 212-call histogram sweep stays in +the WaveScope driver; this table is the mean-Nq next_n × ctx grid. """ from __future__ import annotations @@ -40,8 +45,14 @@ INDEX_DIM = HEAD_DIM + 4 MAX_NN = 8 WIDE_MAX_MODEL_LEN = 1 << 20 +# probe_1079 decode (bs64): 44362 × 64-token pages, output stride 258048. +PROBE_NUM_BLOCKS = 44362 +PROBE_MAX_MODEL_LEN = 258048 +PROBE_MEAN_NQ = 117 +PROBE_SEED = 1079 _E4M3_NATIVE = get_fp8_e4m3_dtype() _Q_DTYPE = {"fn": torch.float8_e4m3fn} +_PROBE_KV = None try: from aiter.ops.triton.attention.pa_mqa_logits import ( @@ -226,6 +237,74 @@ def _build_inputs( ) +def scattered_block_tables( + nq, pages, num_blocks=PROBE_NUM_BLOCKS, max_model_len=PROBE_MAX_MODEL_LEN +): + """Random physical pages; consecutive logical pages are not adjacent in HBM.""" + used = nq * pages + max_blocks_seq = max_model_len // KV_BLOCK_SIZE + if used <= num_blocks: + ids = torch.randperm(num_blocks, device="cuda")[:used] + else: + reps = (used + num_blocks - 1) // num_blocks + ids = torch.cat( + [torch.randperm(num_blocks, device="cuda") for _ in range(reps)] + )[:used] + tables = torch.zeros((nq, max_blocks_seq), dtype=torch.int32, device="cuda") + tables[:, :pages] = ids.reshape(nq, pages).to(torch.int32) + return tables + + +def _probe_kv_pool(): + global _PROBE_KV + if _PROBE_KV is None: + fp8 = get_fp8_e4m3_dtype() + kv = torch.randn( + (PROBE_NUM_BLOCKS, KV_BLOCK_SIZE, 1, HEAD_DIM), dtype=torch.bfloat16 + ) + raw = kv_cache_cast_to_fp8(kv, fp8) + _PROBE_KV = (raw, preshuffle_kv_data(raw, HEAD_DIM), fp8) + return _PROBE_KV + + +def _build_probe_inputs(nq, next_n, kv_len, q_dtype): + """Compact decode tensors matching probe_1079 / the indexer call.""" + if kv_len % KV_BLOCK_SIZE: + raise ValueError(f"kv_len {kv_len} must be a multiple of kvb={KV_BLOCK_SIZE}") + if kv_len > PROBE_MAX_MODEL_LEN: + raise ValueError( + f"kv_len {kv_len} exceeds probe max_model_len {PROBE_MAX_MODEL_LEN}" + ) + torch.manual_seed(PROBE_SEED) + random.seed(PROBE_SEED) + pages = kv_len // KV_BLOCK_SIZE + kv_raw, kv_kernel, fp8_dtype = _probe_kv_pool() + q = torch.randn((nq, next_n, HEADS, HEAD_DIM), dtype=torch.bfloat16) + weights = torch.randn((nq * next_n, HEADS), dtype=torch.float32) + context_lens = torch.full((nq,), kv_len, device="cuda", dtype=torch.int32) + tables = scattered_block_tables(nq, pages) + out = torch.full( + (nq * next_n, PROBE_MAX_MODEL_LEN), + float("-inf"), + device="cuda", + dtype=torch.float32, + ) + return ( + Inputs( + q, + q.to(q_dtype), + kv_raw, + weights, + context_lens, + tables, + PROBE_MAX_MODEL_LEN, + fp8_dtype, + ), + kv_kernel, + out, + ) + + def _kernel_inputs(inp, batch_size, next_n, head_dim): kv_cache_kernel = preshuffle_kv_data(inp.kv_cache_fp8, head_dim) out = torch.full( @@ -414,6 +493,68 @@ def gluon(): return _time_candidates(candidates, ref, flops, nbytes, "paged fp8_mqa_logits") +@benchmark() +def test_fp8_paged_mqa_logits_probe1079(nq, next_n, kv_len, dtype): + """Mean-Nq decode: scattered 44362-block pool, production output stride.""" + q_dtype = _Q_DTYPE[dtype] + inp, kv_cache_kernel, out = _build_probe_inputs(nq, next_n, kv_len, q_dtype) + scored = out[:, :kv_len] + ref = run_torch( + inp.q, + inp.kv_cache_fp8, + inp.weights, + inp.context_lens, + inp.block_tables, + kv_len, + inp.fp8_dtype, + block_size=KV_BLOCK_SIZE, + ) + + def flydsl(): + out.fill_(float("-inf")) + flydsl_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out, + inp.context_lens, + inp.block_tables, + PROBE_MAX_MODEL_LEN, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + ) + return scored + + candidates = {"flydsl": (flydsl, scored)} + if deepgemm_fp8_paged_mqa_logits is not None: + out_g = out.clone() + scored_g = out_g[:, :kv_len] + + def gluon(): + out_g.fill_(float("-inf")) + deepgemm_fp8_paged_mqa_logits( + inp.q_fp8, + kv_cache_kernel, + inp.weights, + out_g, + inp.context_lens, + inp.block_tables, + PROBE_MAX_MODEL_LEN, + ChunkK=256, + Preshuffle=True, + KVBlockSize=KV_BLOCK_SIZE, + WavePerEU=2, + ) + return scored_g + + candidates["gluon"] = (gluon, scored_g) + + flops, nbytes = _roofline( + nq, next_n, kv_len, inp.context_lens, inp.q_fp8, inp.weights, scored + ) + return _time_candidates(candidates, ref, flops, nbytes, "paged probe1079") + + @benchmark() def test_fp8_paged_mqa_logits_ragged(batch, next_n, kv_len, split_kv): inp = _build_inputs( @@ -634,6 +775,29 @@ def main(): default=[65], help="Batch for the wide-output i32-offset table.", ) + parser.add_argument( + "--hist-nq", + type=int, + nargs="*", + default=[PROBE_MEAN_NQ], + help="Nq for the probe_1079 table (mean of the 212-call histogram is 117).\n" + " e.g.: --hist-nq 16 117 256", + ) + parser.add_argument( + "--hist-next-n", + type=int, + nargs="*", + default=list(range(1, MAX_NN + 1)), + help="Compact next_n for the probe_1079 table.\n e.g.: --hist-next-n 1 8", + ) + parser.add_argument( + "--hist-kv-len", + type=int, + nargs="*", + default=[8192], + help="Tokens/seq for the probe_1079 table (scattered 8 KiB pages).\n" + " e.g.: --hist-kv-len 8192 131072", + ) args = parser.parse_args() compact = [ @@ -673,6 +837,15 @@ def main(): if wide: _summarize("fp8_paged_mqa_logits_wide_out", wide) + hist = [ + test_fp8_paged_mqa_logits_probe1079(nq, next_n, kv_len, dtype) + for dtype, nq, next_n, kv_len in itertools.product( + args.dtype, args.hist_nq, args.hist_next_n, args.hist_kv_len + ) + ] + if hist: + _summarize("fp8_paged_mqa_logits_probe1079", hist) + if __name__ == "__main__": main()