diff --git a/aiter/ops/flydsl/__init__.py b/aiter/ops/flydsl/__init__.py index fae2c0be94d..3b6f3a67388 100644 --- a/aiter/ops/flydsl/__init__.py +++ b/aiter/ops/flydsl/__init__.py @@ -54,6 +54,14 @@ ".fp8_mqa_logits_kernels", "flydsl_fp8_mqa_logits", ), + "flydsl_fp8_paged_mqa_logits": ( + ".kernels.mqa_logits.fp8_paged_mqa_logits_gfx950", + "flydsl_fp8_paged_mqa_logits", + ), + "flydsl_fp8_paged_mqa_logits_gfx950": ( + ".kernels.mqa_logits.fp8_paged_mqa_logits_gfx950", + "flydsl_fp8_paged_mqa_logits_gfx950", + ), "flydsl_hgemm": (".gemm_kernels", "flydsl_hgemm"), "flydsl_hstu_attention": ( ".hstu_attention", @@ -116,6 +124,8 @@ "flydsl_flash_attn_fp8_supported", "flydsl_flash_attn_func", "flydsl_fp8_mqa_logits", + "flydsl_fp8_paged_mqa_logits", + "flydsl_fp8_paged_mqa_logits_gfx950", "flydsl_hgemm", "flydsl_hstu_attention", "flydsl_hstu_attention_bwd", 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 new file mode 100644 index 00000000000..3fc1c358894 --- /dev/null +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""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. +""" + +from .fp8_paged_mqa_logits_gfx950 import ( # noqa: F401 + flydsl_fp8_paged_mqa_logits, + flydsl_fp8_paged_mqa_logits_gfx950, +) 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..4bf0ec91081 --- /dev/null +++ b/aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.py @@ -0,0 +1,584 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +"""gfx950 H32/D128/KVB64 ragged paged FP8 MQA-logits kernel. + +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 + +import flydsl.compiler as flyc +import flydsl.expr as fx +import torch +from flydsl.expr import 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 + +Vec = fx.Vector + +MAX_NN = 8 +NUM_HEADS = 32 +HEAD_DIM = 128 +INDEX_DIM = HEAD_DIM + 4 +KV_BLOCK_SIZE = 64 +MFMA_M = MFMA_N = 16 +M_TILES = NUM_HEADS // MFMA_M +DREG = 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 + +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 + +_NEUTRAL_E8M0 = 0x7F7F7F7F +_GFX950 = "gfx950" +DEFAULT_COMPILE_HINTS = {"waves_per_eu": 6, "fast_fp_math": True} + + +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 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 _wait_all(): + rocdl.sched_barrier(0) + rocdl.s_waitcnt(vmcnt=0, lgkmcnt=0) + rocdl.sched_barrier(0) + + +def _guarded(pred, fn): + @flyc.jit + def _do(_pred=pred, _fn=fn): + if _pred: + _fn() + + _do() + + +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 _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 _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_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_score(a_pack, b_pack): + 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): + zero = fx.Float32(0.0) + total = zero + for mi in range_constexpr(M_TILES): + for ii in range_constexpr(DREG): + 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 _make_out_row_view(logits, stride_out, row): + byte = fx.Int64(fx.Uint32(row)) * fx.Int64(fx.Uint32(stride_out)) * 4 + return GTensor( + logits, + dtype=T.f32, + shape=(-1,), + static_bytes_offset_i64=byte, + ) + + +@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 = ( + "fp8_paged_mqa_logits_gfx950_H32_D128_bkv64_kvb64_" + "nn8_w8_lds2_mfma16_ntg2l_flydsl" + ) + + @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, + weights: fx.Tensor, + 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, + 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, + ): + 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) + + 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,)) + + 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) + + 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, + ) + + # 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, + ) + + _guarded(pred, _copy_q) + + w_batch_base = pid_batch * (rows_per_batch * 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, + pid_batch * max_block_len + safe_page, + vec_width=1, + dtype=T.i32, + is_scalar=True, + ) + ) + + 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, + ) + + _guarded(pred, _copy) + + def _store_logit(row, col, value): + 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 + + _guarded( + (lane_div_16 == 0) + & (col < context_len) + & (col <= context_len - nn + row), + _write, + ) + + 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, + ) + + def _score_row(r): + a_packs = [ + _load_q_pack_lds( + q_lds_base, + 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, r, 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( + r, + page * KV_BLOCK_SIZE + token_base + lane_mod_16, + value, + ) + + 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). + _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, + ) + + _guarded( + (wave >= COMPUTE_WAVES) & (next_page < page_hi), + _prefetch, + ) + _wait_all() + gpu.barrier() + yield [next_bank] + + @flyc.jit + def launch( + Q, + KV_cache, + weights, + out_logits, + context_lens, + kv_indices, + next_n_lens, + grid_blocks, + batch_size, + split_kv, + rows_per_batch, + has_next_n_lens, + stride_q_batch, + 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_lens, + batch_size, + split_kv, + rows_per_batch, + has_next_n_lens, + stride_q_batch, + max_block_len, + stride_out, + ).launch( + grid=(fx.Int64(grid_blocks), 1, 1), + block=(THREADS, 1, 1), + stream=stream, + ) + + launch.compile_hints = dict(DEFAULT_COMPILE_HINTS) + launch.kernel_name = kernel_name + return launch + + +@lru_cache(maxsize=1) +def _compile(): + return _build_kernel() + + +def flydsl_fp8_paged_mqa_logits( + q_fp8, + kv_cache, + weights, + out_logits, + context_lens, + 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 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, 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, 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(): + raise ValueError(f"q_fp8 must be native FP8 E4M3, got {q_fp8.dtype}") + 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, 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 * 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 = 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") + + # 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, + (int(context_lens.max().item()) + KV_BLOCK_SIZE - 1) // KV_BLOCK_SIZE, + ) + 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() + 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( + launcher, + q_fp8, + kv_cache.reshape(-1), + weights, + out_logits, + context_lens, + kv_indices, + next_n_lens, + 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)), + stream, + ) + return out_logits + + +flydsl_fp8_paged_mqa_logits_gfx950 = flydsl_fp8_paged_mqa_logits 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..c638d8d6f76 --- /dev/null +++ b/op_tests/test_flydsl_fp8_paged_mqa_logits.py @@ -0,0 +1,851 @@ +# 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. + +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) + +``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 + +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 +# 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 ( + 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 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( + (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, dtype): + q_dtype = _Q_DTYPE[dtype] + inp = _build_inputs( + 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( + 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_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( + 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")) + assert torch.all(out[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, + ) + 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(): + wide.fill_(float("-inf")) + 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, + ) + return scored + + flops, nbytes = _roofline( + 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" + ) + + +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"], + choices=list(_Q_DTYPE), + 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=[1, 2, 4, 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=[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, 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", + type=int, + nargs="*", + 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 = [ + 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) + + 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()