diff --git a/.github/scripts/split_tests.sh b/.github/scripts/split_tests.sh index 56f5716185..451dc94199 100755 --- a/.github/scripts/split_tests.sh +++ b/.github/scripts/split_tests.sh @@ -103,7 +103,7 @@ if [[ "$TEST_TYPE" == "aiter" ]]; then FILE_TIMES[op_tests/test_batched_gemm_a8w8.py]=51 FILE_TIMES[op_tests/test_mha_varlen_large_kv.py]=46 FILE_TIMES[op_tests/test_mla_reduce.py]=44 - FILE_TIMES[op_tests/test_pa_sparse_prefill_opus.py]=44 + FILE_TIMES[op_tests/test_pa_sparse_prefill.py]=44 FILE_TIMES[op_tests/test_pa_ragged.py]=40 FILE_TIMES[op_tests/test_moeTopkSoftmax.py]=39 FILE_TIMES[op_tests/test_moe_sorting_mxfp4.py]=39 diff --git a/aiter/__init__.py b/aiter/__init__.py index 5aac376c0d..c6eaa03f37 100644 --- a/aiter/__init__.py +++ b/aiter/__init__.py @@ -112,6 +112,7 @@ def getLogger(): from .ops.moe_sorting import * from .ops.moe_sorting_opus import * from .ops.moe_mxfp4_aux import * + from .ops.mla_sparse_prefill import * from .ops.pa_sparse_prefill_opus import * from .ops.pos_encoding import * from .ops.cache import * diff --git a/aiter/jit/optCompilerConfig.json b/aiter/jit/optCompilerConfig.json index ee2e6347a9..7699707f68 100644 --- a/aiter/jit/optCompilerConfig.json +++ b/aiter/jit/optCompilerConfig.json @@ -103,6 +103,17 @@ "verbose": "False", "blob_gen_cmd": "f'{AITER_META_DIR}/hsa/codegen.py -m mla_v4 --output_dir {{}}'" }, + "module_mla_sparse_prefill_asm": { + "srcs": [ + "f'{AITER_CSRC_DIR}/py_itfs_cu/asm_mla_sparse_prefill.cu'" + ], + "flags_extra_cc": [], + "flags_extra_hip": [], + "extra_ldflags": "None", + "extra_include": [], + "verbose": "False", + "blob_gen_cmd": "f'{AITER_META_DIR}/hsa/codegen.py -m mla_v4 --output_dir {{}}'" + }, "module_cache": { "srcs": [ "f'{AITER_CSRC_DIR}/pybind/cache_pybind.cu'", diff --git a/aiter/ops/mla_sparse_prefill.py b/aiter/ops/mla_sparse_prefill.py new file mode 100644 index 0000000000..bcb94f056a --- /dev/null +++ b/aiter/ops/mla_sparse_prefill.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: MIT +# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. + +import torch + +from ..jit.core import compile_ops +from ..jit.utils.chip_info import get_gfx_runtime +from ..jit.utils.torch_guard import torch_compile_guard + +MD_NAME = "module_mla_sparse_prefill_asm" + + +# NOTE: ctypes binds positionally off this signature -- the argument order here +# must match `mla_sparse_prefill_fp8_asm_fwd` in +# csrc/py_itfs_cu/asm_mla_sparse_prefill.cu one for one. A silent reorder here +# is a silent wrong-pointer launch, not a compile error. +@compile_ops(MD_NAME, fc_name="mla_sparse_prefill_fp8_asm_fwd", ffi_type="ctypes") +def mla_sparse_prefill_fp8_asm_fwd( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + unified_kv_nope: torch.Tensor, + unified_kv_rope: torch.Tensor, + kv_indices_prefix: torch.Tensor, + kv_indptr_prefix: torch.Tensor, + kv_nope: torch.Tensor, + kv_rope: torch.Tensor, + kv_indices_extend: torch.Tensor, + kv_indptr_extend: torch.Tensor, + attn_sink: torch.Tensor, + out: torch.Tensor, + softmax_scale: float, +) -> None: ... + + +def _mla_sparse_prefill_fp8_asm_fake( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + unified_kv_nope: torch.Tensor, + unified_kv_rope: torch.Tensor, + kv_indices_prefix: torch.Tensor, + kv_indptr_prefix: torch.Tensor, + kv_nope: torch.Tensor, + kv_rope: torch.Tensor, + kv_indices_extend: torch.Tensor, + kv_indptr_extend: torch.Tensor, + attn_sink: torch.Tensor, + softmax_scale: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + if out is not None: + return out + t, h, _ = q_nope.shape + return torch.empty((t, h, 512), dtype=torch.bfloat16, device=q_nope.device) + + +@torch_compile_guard(mutates_args=["out"], gen_fake=_mla_sparse_prefill_fp8_asm_fake) +def mla_sparse_prefill_fp8_asm( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + unified_kv_nope: torch.Tensor, + unified_kv_rope: torch.Tensor, + kv_indices_prefix: torch.Tensor, + kv_indptr_prefix: torch.Tensor, + kv_nope: torch.Tensor, + kv_rope: torch.Tensor, + kv_indices_extend: torch.Tensor, + kv_indptr_extend: torch.Tensor, + attn_sink: torch.Tensor, + softmax_scale: float, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """Sparse prefill attention with split fp8 NoPE and bf16 RoPE inputs (asm). + + Signature-compatible with + :func:`aiter.ops.pa_sparse_prefill_opus.pa_sparse_prefill_fp8_opus`. + + Args: + q_nope: ``[T, H, 512]`` fp8 query without positional encoding. + q_rope: ``[T, H, 64]`` bf16 query RoPE encoding part. + unified_kv_nope: ``[total_pages, 512]`` fp8 prefix KV NoPE source. + unified_kv_rope: ``[total_pages, 64]`` bf16 prefix KV RoPE source. + kv_indices_prefix: ``[total_prefix]`` int32 row indices into the prefix + sources, concatenated per token. + kv_indptr_prefix: ``[T+1]`` int32 CSR row pointers. + kv_nope: ``[total_tokens, 512]`` fp8 extend KV NoPE source. + kv_rope: ``[total_tokens, 64]`` bf16 extend KV RoPE source. + kv_indices_extend: ``[total_extend]`` int32 row indices into the extend + sources, concatenated per token. + kv_indptr_extend: ``[T+1]`` int32 CSR row pointers. + attn_sink: ``[H]`` fp32 per-head softmax-denom bias. + softmax_scale: float scalar applied to the combined QK^T scores. + out: Optional ``[T, H, 512]`` bf16 output buffer; allocated + if ``None``. + + Returns: + ``out`` (``[T, H, 512]`` bf16). + """ + gfx = get_gfx_runtime() + if gfx != "gfx1250": + raise RuntimeError(f"mla_sparse_prefill_fp8_asm requires gfx1250, got {gfx}") + + if q_nope.dtype != unified_kv_nope.dtype or q_nope.dtype != kv_nope.dtype: + raise RuntimeError( + f"NoPE dtype mismatch: q_nope={q_nope.dtype}, " + f"unified_kv_nope={unified_kv_nope.dtype}, kv_nope={kv_nope.dtype}" + ) + if q_rope.dtype != torch.bfloat16: + raise RuntimeError(f"q_rope must be bf16, got {q_rope.dtype}") + + t, h = q_nope.shape[0], q_nope.shape[1] + if h != 128: + # Hard constraint, not a dispatch miss. + raise RuntimeError(f"mla_sparse_prefill_fp8_asm requires H == 128, got {h}") + if out is None: + out = torch.empty((t, h, 512), dtype=torch.bfloat16, device=q_nope.device) + elif out.shape != (t, h, 512) or out.dtype != torch.bfloat16: + raise RuntimeError( + f"out shape/dtype mismatch: got shape={tuple(out.shape)} dtype={out.dtype}, " + f"expected shape={(t, h, 512)} dtype={torch.bfloat16}" + ) + + mla_sparse_prefill_fp8_asm_fwd( + q_nope, + q_rope, + unified_kv_nope, + unified_kv_rope, + kv_indices_prefix, + kv_indptr_prefix, + kv_nope, + kv_rope, + kv_indices_extend, + kv_indptr_extend, + attn_sink, + out, + float(softmax_scale), + ) + return out + + +__all__ = [ + "mla_sparse_prefill_fp8_asm", + "mla_sparse_prefill_fp8_asm_fwd", +] diff --git a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py index 80af195c12..bd86c0f3d8 100644 --- a/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py +++ b/aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py @@ -291,7 +291,12 @@ def _sparse_attn_prefill_kernel( BLOCK_D: tl.constexpr, BLOCK_K: tl.constexpr, ): - query_idx = tl.program_id(0) + # 64-bit before the multiply, same reasoning as `slot_off` below: the + # program id fits 32 bits, but `query_idx * q_stride_t` does not once + # num_queries passes 32K, because q_stride_t is num_heads * head_dim + # (128 * 512 = 65536) in the V4 layout. Unlike the pool read, the wrapped + # offset lands outside the q/out allocations, so it page-faults. + query_idx = tl.program_id(0).to(tl.int64) pid_h = tl.program_id(1) head_offsets = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) diff --git a/csrc/py_itfs_cu/asm_mla_sparse_prefill.cu b/csrc/py_itfs_cu/asm_mla_sparse_prefill.cu new file mode 100644 index 0000000000..790f916ef0 --- /dev/null +++ b/csrc/py_itfs_cu/asm_mla_sparse_prefill.cu @@ -0,0 +1,260 @@ +// SPDX-License-Identifier: MIT +// Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. +// +// gfx1250 (MI400) DSA sparse-prefill MLA asm dispatcher. + +#include "aiter_tensor.h" +#include "aiter_ctypes_error.h" +#include "asm_mla_v4_configs.hpp" +#include +#include +#include +#include +#include + +AITER_CTYPES_ERROR_DEF + +namespace { + +struct __attribute__((packed)) MlaSparsePrefillKargs +{ + const void* q_nope; // 0x00 [T, H, 512] fp8 (448 data + 14 e8m0 scales + pad) + const void* q_rope; // 0x08 [T, H, 64] bf16 + const void* unified_kv_nope; // 0x10 [total_pages, 512] fp8 <- prefix source + const void* unified_kv_rope; // 0x18 [total_pages, 64] bf16 + const void* kv_nope; // 0x20 [total_tokens, 512] fp8 <- extend source + const void* kv_rope; // 0x28 [total_tokens, 64] bf16 + const void* attn_sink; // 0x30 [H] fp32 + void* out; // 0x38 [T, H, 512] bf16 (read_write) + const void* kv_indptr_prefix; // 0x40 [T+1] int32 + const void* kv_indices_prefix; // 0x48 [nnz_prefix] int32 + const void* kv_indptr_extend; // 0x50 [T+1] int32 + const void* kv_indices_extend; // 0x58 [nnz_extend] int32 + float softmax_scale; // 0x60 f32 by_value + unsigned int _tail_pad; // 0x64 +}; + +static_assert(sizeof(MlaSparsePrefillKargs) == 104, + "kernarg packet must stay 104 bytes (.kernarg_segment_size in the .co)"); +static_assert(offsetof(MlaSparsePrefillKargs, q_nope) == 0x00, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, q_rope) == 0x08, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, unified_kv_nope) == 0x10, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, unified_kv_rope) == 0x18, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, kv_nope) == 0x20, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, kv_rope) == 0x28, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, attn_sink) == 0x30, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, out) == 0x38, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, kv_indptr_prefix) == 0x40, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, kv_indices_prefix) == 0x48, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, kv_indptr_extend) == 0x50, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, kv_indices_extend) == 0x58, "kernarg offset drift"); +static_assert(offsetof(MlaSparsePrefillKargs, softmax_scale) == 0x60, "kernarg offset drift"); + +constexpr int kHeads = 128; // == the CSV's Gqa for the one prefill row +constexpr int kHeadDim = 512; // NoPE packed row: 448 fp8 + 14 e8m0 scales + 50 pad +constexpr int kRopeDim = 64; +constexpr int kBlockDim = 128; // 4 x wave32 + +std::string dtype_name(const aiter_tensor_t* t) +{ + return AiterDtype_to_str(t->dtype()); +} + +// `allow_empty` covers the CSR index lists: a region with no entries at all is +// a legal input (the kernel branches straight past an empty prefix stream and +// parks the prefix->extend switch index out of reach when extend is empty), and +// torch hands us a NULL data pointer for a zero-element tensor. Everything else +// must be a real allocation. +void check_tensor(const aiter_tensor_t* t, + const char* name, + AiterDtype want, + bool allow_empty = false) +{ + AITER_CHECK(t != nullptr, "mla_sparse_prefill_fp8_asm: `", name, "` must not be NULL"); + AITER_CHECK(t->data_ptr() != nullptr || (allow_empty && t->numel() == 0), + "mla_sparse_prefill_fp8_asm: `", name, "` has a NULL data pointer"); + AITER_CHECK(t->dtype() == want, + "mla_sparse_prefill_fp8_asm: `", name, "` must be ", + AiterDtype_to_str(want), ", got ", dtype_name(t)); + AITER_CHECK(t->is_contiguous(), + "mla_sparse_prefill_fp8_asm: `", name, "` must be contiguous"); +} + +// CSR row pointers are read as [T+1] int32. An empty index list is legal. +void check_csr(const aiter_tensor_t* indptr, + const aiter_tensor_t* indices, + const char* indptr_name, + const char* indices_name, + int64_t t) +{ + check_tensor(indptr, indptr_name, AITER_DTYPE_i32); + check_tensor(indices, indices_name, AITER_DTYPE_i32, /*allow_empty=*/true); + AITER_CHECK(indptr->numel() >= static_cast(t + 1), + "mla_sparse_prefill_fp8_asm: `", indptr_name, "` must have at least T+1 = ", + t + 1, " elements, got ", indptr->numel()); +} + +} // namespace + +AITER_CTYPES_DEFINE_ENTRYPOINT_VOID( + mla_sparse_prefill_fp8_asm_fwd, + (aiter_tensor_t * q_nope, // [T, H, 512] fp8 + aiter_tensor_t* q_rope, // [T, H, 64] bf16 + aiter_tensor_t* unified_kv_nope, // [total_pages, 512] fp8 + aiter_tensor_t* unified_kv_rope, // [total_pages, 64] bf16 + aiter_tensor_t* kv_indices_prefix, // [nnz_prefix] int32 + aiter_tensor_t* kv_indptr_prefix, // [T+1] int32 + aiter_tensor_t* kv_nope, // [total_tokens, 512] fp8 + aiter_tensor_t* kv_rope, // [total_tokens, 64] bf16 + aiter_tensor_t* kv_indices_extend, // [nnz_extend] int32 + aiter_tensor_t* kv_indptr_extend, // [T+1] int32 + aiter_tensor_t* attn_sink, // [H] fp32 + aiter_tensor_t* out, // [T, H, 512] bf16 (written) + float softmax_scale, + hipStream_t stream), + (q_nope, + q_rope, + unified_kv_nope, + unified_kv_rope, + kv_indices_prefix, + kv_indptr_prefix, + kv_nope, + kv_rope, + kv_indices_extend, + kv_indptr_extend, + attn_sink, + out, + softmax_scale, + stream)) +{ + const std::string arch_id = get_gpu_arch(); + AITER_CHECK(arch_id == "gfx1250", + "mla_sparse_prefill_fp8_asm: only gfx1250 is supported, got ", arch_id); + + check_tensor(q_nope, "q_nope", AITER_DTYPE_fp8); + check_tensor(q_rope, "q_rope", AITER_DTYPE_bf16); + check_tensor(unified_kv_nope, "unified_kv_nope", AITER_DTYPE_fp8); + check_tensor(unified_kv_rope, "unified_kv_rope", AITER_DTYPE_bf16); + check_tensor(kv_nope, "kv_nope", AITER_DTYPE_fp8); + check_tensor(kv_rope, "kv_rope", AITER_DTYPE_bf16); + check_tensor(attn_sink, "attn_sink", AITER_DTYPE_fp32); + check_tensor(out, "out", AITER_DTYPE_bf16); + + AITER_CHECK(q_nope->dim() == 3, + "mla_sparse_prefill_fp8_asm: `q_nope` must be 3-D [T, H, 512], got ndim=", + q_nope->dim()); + const int64_t t = q_nope->size(0); + const int64_t h = q_nope->size(1); + + AITER_CHECK(h == kHeads, + "mla_sparse_prefill_fp8_asm: this kernel is built for exactly H=", kHeads, + " heads (one workgroup serves one query token x ", kHeads, + " heads, and the Q address math requires gridDim.y == 1); got H=", h); + AITER_CHECK(q_nope->size(2) == kHeadDim, + "mla_sparse_prefill_fp8_asm: `q_nope` last dim must be ", kHeadDim, ", got ", + q_nope->size(2)); + AITER_CHECK(q_rope->size(2) == kRopeDim, + "mla_sparse_prefill_fp8_asm: `q_rope` last dim must be ", kRopeDim, ", got ", + q_rope->size(2)); + AITER_CHECK(unified_kv_nope->size(-1) == kHeadDim && kv_nope->size(-1) == kHeadDim, + "mla_sparse_prefill_fp8_asm: KV NoPE rows must be ", kHeadDim, " wide, got ", + unified_kv_nope->size(-1), " / ", kv_nope->size(-1)); + AITER_CHECK(unified_kv_rope->size(-1) == kRopeDim && kv_rope->size(-1) == kRopeDim, + "mla_sparse_prefill_fp8_asm: KV RoPE rows must be ", kRopeDim, " wide, got ", + unified_kv_rope->size(-1), " / ", kv_rope->size(-1)); + AITER_CHECK(attn_sink->numel() == static_cast(h), + "mla_sparse_prefill_fp8_asm: `attn_sink` must have H=", h, " elements, got ", + attn_sink->numel()); + AITER_CHECK(out->dim() == 3 && out->size(0) == t && out->size(1) == h && + out->size(2) == kHeadDim, + "mla_sparse_prefill_fp8_asm: `out` must be [", t, ", ", h, ", ", kHeadDim, "]"); + + check_csr(kv_indptr_prefix, kv_indices_prefix, "kv_indptr_prefix", "kv_indices_prefix", t); + check_csr(kv_indptr_extend, kv_indices_extend, "kv_indptr_extend", "kv_indices_extend", t); + + MlaSparsePrefillKargs args; + size_t arg_size = sizeof(args); + args.q_nope = q_nope->data_ptr(); + args.q_rope = q_rope->data_ptr(); + args.unified_kv_nope = unified_kv_nope->data_ptr(); + args.unified_kv_rope = unified_kv_rope->data_ptr(); + args.kv_nope = kv_nope->data_ptr(); + args.kv_rope = kv_rope->data_ptr(); + args.attn_sink = attn_sink->data_ptr(); + args.out = out->data_ptr(); + args.kv_indptr_prefix = kv_indptr_prefix->data_ptr(); + args.kv_indices_prefix = kv_indices_prefix->data_ptr(); + args.kv_indptr_extend = kv_indptr_extend->data_ptr(); + args.kv_indices_extend = kv_indices_extend->data_ptr(); + args.softmax_scale = softmax_scale; + args._tail_pad = 0; + + CFG* config_map = &cfg_mla_v4_asm; + static SynchronizedCache impl_ptr_map; + AiterAsmKernel* impl_ptr = nullptr; + + // The manifest is shared with the v4 sparse DECODE kernels, so prefill must be + // part of the key -- otherwise a decode row with a matching Gqa would be picked + // and launched with the wrong kernarg packet. Gqa is the head count (one + // workgroup serves one query token x Gqa heads) and qSeqLen is 1 by + // construction on this path. + std::string kernelName; + for(const auto& el : *config_map) + { + if(el.first.find(arch_id) != 0) + continue; + const auto& cfg = el.second; + if(cfg.prefill != 1 || cfg.causal != 0 || cfg.lse != 0 || cfg.ps != 0) + continue; + if(cfg.qType != "fp8" || cfg.kvType != "fp8") + continue; + if(cfg.Gqa != static_cast(h) || cfg.qSeqLen != 1) + continue; + kernelName = el.first; + break; + } + AITER_CHECK(!kernelName.empty(), + "mla_sparse_prefill_fp8_asm: no prefill kernel for arch=", arch_id, + " qType=fp8 kvType=fp8 Gqa=", h); + + auto it = config_map->find(kernelName); + AITER_CHECK(it != config_map->end(), + "mla_sparse_prefill_fp8_asm: kernel not found: ", kernelName); + { + const auto& cfg = it->second; + const char* name = cfg.knl_name.c_str(); + const char* co_name = cfg.co_name.c_str(); + impl_ptr = + &impl_ptr_map.get_or_create(name, [&]() { return AiterAsmKernel(name, co_name); }); + } + + // gdx = one workgroup per query token. gdy MUST stay 1 (see kHeads note). + AITER_CHECK(t >= 0 && (t >> 31) == 0, "mla_sparse_prefill_fp8_asm: T too large: ", t); + const int gdx = static_cast(t); + if(gdx == 0) + return; // nothing to do; `out` stays as the caller left it + + if(const char* dbg = std::getenv("AITER_MLA_SPARSE_PREFILL_DUMP_KERNARG")) + { + if(dbg[0] == '1') + { + fprintf(stderr, "[aiter pa_sparse_prefill kernarg %zuB]\n", arg_size); + const uint8_t* bytes = reinterpret_cast(&args); + for(size_t i = 0; i < arg_size; ++i) + fprintf(stderr, "%02x%s", bytes[i], ((i + 1) % 16 == 0) ? "\n" : " "); + fprintf(stderr, "\n[aiter grid (%d,1,1) block (%d,1,1)]\n", gdx, kBlockDim); + fflush(stderr); + } + } + + const HipDeviceGuard device_guard(q_nope->device_id); + impl_ptr->launch_kernel({&args, + &arg_size, + gdx, // gdx: one query token per workgroup + 1, // gdy: MUST be 1 + 1, // gdz + kBlockDim, // bdx: 4 x wave32 + 1, // bdy + 1, // bdz + stream}); +} diff --git a/hsa/gfx1250/mla_v4/mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co b/hsa/gfx1250/mla_v4/mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co new file mode 100755 index 0000000000..8198564688 Binary files /dev/null and b/hsa/gfx1250/mla_v4/mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co differ diff --git a/hsa/gfx1250/mla_v4/mla_v4_asm.csv b/hsa/gfx1250/mla_v4/mla_v4_asm.csv index 1fe4ba4d34..a2607c25dc 100644 --- a/hsa/gfx1250/mla_v4/mla_v4_asm.csv +++ b/hsa/gfx1250/mla_v4/mla_v4_asm.csv @@ -1,3 +1,14 @@ +// +// prefill=1 rows: DSA sparse PREFILL (not decode, and despite the name not causal +// either -- there is no qo_len/kv_len/causal mask in the kernel; causality is +// pre-baked by the caller into two per-query-token CSR index lists). One +// workgroup serves one query token x Gqa heads, so Gqa is a HARD constraint, +// not a dispatch hint. These use their own compact 104-byte kernarg ABI +// (csrc/py_itfs_cu/asm_mla_sparse_prefill.cu), not the decode packet. +// qh128 32mx4_32nx1_sparse_pfl <- my_agent_skills/0805_base/32mx4_32nx1.sp3 +// (the .sp3 keeps its tile-shape name; the .co/symbol follow the aiter +// naming convention here) qType,kvType,Gqa,ps,qSeqLen,prefill,causal,lse,knl_name,co_name fp8,fp8,64,0,1,0,0,0,_ZN5aiter35mla_a8w8_qh64_1tg_16mx4_64nx1_sparseE,mla_a8w8_qh64_1tg_16mx4_64nx1_sparse.co fp8,fp8,16,0,4,0,0,0,_ZN5aiter35mla_a8w8_qh16_1tg_16mx4_64nx1_sparseE,mla_a8w8_qh16_1tg_16mx4_64nx1_sparse.co +fp8,fp8,128,0,1,1,0,0,_ZN5aiter40mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pflE,mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co diff --git a/op_tests/test_pa_sparse_prefill_opus.py b/op_tests/test_pa_sparse_prefill.py similarity index 55% rename from op_tests/test_pa_sparse_prefill_opus.py rename to op_tests/test_pa_sparse_prefill.py index e1530cc252..866638142d 100644 --- a/op_tests/test_pa_sparse_prefill_opus.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -1,53 +1,12 @@ # SPDX-License-Identifier: MIT # Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved. -"""Tests for aiter's OPUS-based sparse paged prefill attention. +"""Correctness and performance tests for sparse prefill attention. -We validate both precision variants of the OPUS two-region sparse paged -prefill kernel (gfx950 compiled from source, gfx1250 loaded from a prebuilt -code object) against explicit PyTorch references (per-token online-softmax + -per-head sink): - -* ``pa_sparse_prefill_opus`` -- bf16/fp16 Q/K/V/O in a single ``D=512`` - head-dim tensor. -* ``pa_sparse_prefill_fp8_opus`` -- split-precision DSA inputs: NoPE part - in fp8 (with embedded per-32-block E8M0 scales, 448 + 14 scales + pad = - 512 fp8 slots/row) plus a bf16 RoPE part (64), bf16 output. - -A single ``prec`` axis (``"bf16"`` / ``"fp16"`` / ``"fp8"``) drives both -paths through the same shape/mode sweep, so fp8 is exercised with the same -test parameters as bf16/fp16. - -The same harness drives: - -* A pytest-parametrised correctness sweep (CI). -* A standalone CLI for single-point runs with optional benchmark and - TFLOPS reporting, mirroring the style of ``op_tests/test_batch_prefill.py``. - -Example CLI usage (inside the aiter source tree):: - - # default sweep: N x H_Q x total_pages x {sparse, dense} x {bf16, fp8}, with bench - PYTHONPATH=. python3 op_tests/test_pa_sparse_prefill_opus.py - - # only dense CSR for both prefix and extend - PYTHONPATH=. python3 op_tests/test_pa_sparse_prefill_opus.py --mode dense - - # single shape, fp8 only, no correctness check - PYTHONPATH=. python3 op_tests/test_pa_sparse_prefill_opus.py \\ - -n 1024 --h_q 128 --prec fp8 --mode sparse --no-verify - -Reference design notes (Q&A): - -* Why does the reference cast q/k/v to fp32? The GPU kernel accumulates in - fp32 (``D_ACC = float``). Computing the reference in bf16/fp16 would - let softmax / accumulation errors dominate the comparison; fp32 inside - the ref keeps the diff focused on the kernel's intermediate casts and - is the same convention used by ``test_batch_prefill.py``, - ``aiter/test_mha_common.py::attention_ref(upcast=True)`` etc. -* Why are ``kv_indptr`` / ``kv_indices`` cast to int64 in the ref? Only - for the PyTorch ``index_select`` call -- the API requires Long. The - kernel's ABI still consumes int32 indices on the GPU side; the cast - happens entirely on CPU/GPU but on the ref path only. +Each case uses one shared Q/KV/CSR input and compares the applicable +implementations: OPUS (bf16/fp16 or split fp8), gfx1250 MLA asm (split fp8), +and the portable Triton kernel. Pytest runs correctness coverage; the CLI also +supports benchmark and explicit CSR-boundary sweeps. """ from __future__ import annotations @@ -56,17 +15,21 @@ import itertools import math import os -import sys import pandas as pd import pytest import torch +import triton import aiter # noqa: F401 (registers the top-level export) +from aiter.ops.mla_sparse_prefill import mla_sparse_prefill_fp8_asm from aiter.ops.pa_sparse_prefill_opus import ( pa_sparse_prefill_fp8_opus, pa_sparse_prefill_opus, ) +from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4 import ( + _sparse_attn_prefill_kernel, +) from aiter.test_common import benchmark, checkAllclose, perftest # --------------------------------------------------------------------------- @@ -191,7 +154,6 @@ def _ref_pa_sparse_prefill_opus( _FP8_NBLK = _FP8_D_NOPE // 32 # 14 _FP8_BLK = 32 _FP8_MAX = 448.0 # e4m3fn max normal -_FP8_KV_TILE_SIZE = 64 # KV_TILE_SIZE of the fp8 16mx1_16nx4 kernel def _quantize_nope(real: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: @@ -267,29 +229,31 @@ def _ref_pa_sparse_prefill_fp8( # CSR index generators # --------------------------------------------------------------------------- -# Must match the KV_TILE_SIZE template default in -# csrc/include/pa_sparse_prefill_opus.h. The kernel inner loop advances the -# K/V dimension in chunks of this size, so the trailing-tile branches (full / -# half / over-tile) are most likely to break when nnz_per_row sits at one of -# these boundary values. -_KV_TILE_SIZE = 32 +# Cover the KV tile sizes used by the OPUS and asm candidates. The kernel inner +# loop advances the K/V dimension in chunks of these sizes, so the trailing-tile +# branches are most likely to break when nnz_per_row sits at one of these +# boundary values. +_CSR_TILE_SIZES = (32, 64) -def _boundary_nnz(kv_tile_size: int, total_rows: int) -> list: +def _boundary_nnz(kv_tile_sizes, total_rows: int) -> list: """Tile-boundary nnz values seeded into the leading rows of a sparse CSR, mirroring gcnasm/opus_attn/sparse_paged_attn/pa_host.cc:: init_sparse_kv_indices. Clamped into [0, total_rows].""" - cands = [ - 0, - 1, - kv_tile_size - 1, - kv_tile_size, - kv_tile_size + 1, - 2 * kv_tile_size, - 2 * kv_tile_size + 1, - total_rows, - ] - return [max(0, min(v, total_rows)) for v in cands] + if isinstance(kv_tile_sizes, int): + kv_tile_sizes = (kv_tile_sizes,) + cands = {0, 1, total_rows} + for tile_size in kv_tile_sizes: + cands.update( + ( + tile_size - 1, + tile_size, + tile_size + 1, + 2 * tile_size, + 2 * tile_size + 1, + ) + ) + return [max(0, min(v, total_rows)) for v in sorted(cands)] def _random_csr( @@ -297,7 +261,7 @@ def _random_csr( total_rows: int, *, allow_empty: bool = True, - kv_tile_size: int = _KV_TILE_SIZE, + kv_tile_size=_CSR_TILE_SIZES, device: torch.device, seed: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: @@ -368,134 +332,181 @@ def _empty_csr(n: int, *, device: torch.device) -> tuple[torch.Tensor, torch.Ten _MODES = ("sparse", "dense", "empty") -def _make_inputs( +def _fixed_csr(n: int, nnz: int, total_rows: int, *, device): + """CSR with exactly ``nnz`` entries on every row.""" + if nnz == 0: + return _empty_csr(n, device=device) + indptr = torch.arange(n + 1, dtype=torch.int32, device=device) * nnz + indices = ( + torch.arange(nnz, dtype=torch.int32, device=device) % max(total_rows, 1) + ).repeat(n) + return indptr, indices + + +def _make_sparse_case( n: int, h: int, d: int, total_pages: int, total_tokens: int, - dtype: torch.dtype, *, mode: str = "sparse", + nnz_prefix: int | None = None, + nnz_extend: int | None = None, device: torch.device | str = "cuda", seed: int = 0, ) -> dict: - assert mode in _MODES - torch.manual_seed(seed) + """Generate one logical sparse-attention problem shared by all candidates.""" + assert mode in _MODES or mode == "fixed" device = torch.device(device) - - q = (torch.randn(n, h, d, device=device, dtype=torch.float32) * 0.5).to(dtype) - unified_kv = ( - torch.randn(total_pages, d, device=device, dtype=torch.float32) * 0.5 - ).to(dtype) - kv = (torch.randn(total_tokens, d, device=device, dtype=torch.float32) * 0.5).to( - dtype - ) - attn_sink = torch.randn(h, device=device, dtype=torch.float32) * 0.25 - - def _csr(total_rows: int, seed_offset: int): + gen = torch.Generator(device=device) + gen.manual_seed(seed) + q_fp32 = torch.randn(n, h, d, device=device, generator=gen) * 0.5 + unified_kv_fp32 = torch.randn(total_pages, d, device=device, generator=gen) * 0.5 + kv_fp32 = torch.randn(total_tokens, d, device=device, generator=gen) * 0.5 + attn_sink = torch.randn(h, device=device, generator=gen) * 0.25 + + def _csr(total_rows: int, nnz: int | None, seed_offset: int): + if nnz is not None: + return _fixed_csr(n, nnz, total_rows, device=device) if mode == "sparse": return _random_csr( n, total_rows, device=device, + kv_tile_size=_CSR_TILE_SIZES, seed=seed * 2 + seed_offset, ) if mode == "dense": return _dense_csr(n, total_rows, device=device) return _empty_csr(n, device=device) - ip_p, ix_p = _csr(total_pages, 1) - ip_e, ix_e = _csr(total_tokens, 2) - + ip_p, ix_p = _csr(total_pages, nnz_prefix, 1) + ip_e, ix_e = _csr(total_tokens, nnz_extend, 2) return { - "q": q, - "unified_kv": unified_kv, + "q_fp32": q_fp32, + "unified_kv_fp32": unified_kv_fp32, + "kv_fp32": kv_fp32, "kv_indices_prefix": ix_p, "kv_indptr_prefix": ip_p, - "kv": kv, "kv_indices_extend": ix_e, "kv_indptr_extend": ip_e, "attn_sink": attn_sink, } -def _make_inputs_fp8( - n: int, - h: int, - total_pages: int, - total_tokens: int, - *, - mode: str = "sparse", - device: torch.device | str = "cuda", - seed: int = 0, -) -> dict: - """Build split NoPE-fp8 / RoPE-bf16 inputs plus the matching fp32 - reference rows. Returns ``{"kernel": ..., "ref": ...}`` where ``kernel`` - holds the tensors passed to ``pa_sparse_prefill_fp8_opus`` and ``ref`` - holds the dequantized ``*_fp32`` rows passed to ``_ref_pa_sparse_prefill_fp8``. - """ - assert mode in _MODES - torch.manual_seed(seed) - device = torch.device(device) +def _make_single_inputs(case: dict, dtype: torch.dtype) -> dict: + return { + "q": case["q_fp32"].to(dtype), + "unified_kv": case["unified_kv_fp32"].to(dtype), + "kv_indices_prefix": case["kv_indices_prefix"], + "kv_indptr_prefix": case["kv_indptr_prefix"], + "kv": case["kv_fp32"].to(dtype), + "kv_indices_extend": case["kv_indices_extend"], + "kv_indptr_extend": case["kv_indptr_extend"], + "attn_sink": case["attn_sink"], + } - def _streams(rows: int): - nope_fp8, deq = _quantize_nope( - torch.randn(rows, _FP8_D_NOPE, device=device) * 0.5 - ) - rope = (torch.randn(rows, _FP8_D_ROPE, device=device) * 0.5).to(torch.bfloat16) - row_fp32 = torch.cat([deq, rope.to(torch.float32)], dim=1) # [rows, 512] - return nope_fp8, rope, row_fp32 - qn, qr, q_fp32 = _streams(n * h) - qn = qn.reshape(n, h, _FP8_D_NOPE_PADDED) - qr = qr.reshape(n, h, _FP8_D_ROPE) - q_fp32 = q_fp32.reshape(n, h, _FP8_D_HEAD) - ukn, ukr, ukv_fp32 = _streams(total_pages) - kn, kr, kv_fp32 = _streams(total_tokens) +def _make_split_inputs(case: dict) -> dict: + def split_rows(rows: torch.Tensor): + flat = rows.reshape(-1, _FP8_D_HEAD) + nope, deq = _quantize_nope(flat[:, :_FP8_D_NOPE]) + rope = flat[:, _FP8_D_NOPE:].to(torch.bfloat16) + return nope, rope, torch.cat([deq, rope.to(torch.float32)], dim=1) - attn_sink = torch.randn(h, device=device, dtype=torch.float32) * 0.25 + qn, qr, q_fp32 = split_rows(case["q_fp32"]) + ukn, ukr, ukv_fp32 = split_rows(case["unified_kv_fp32"]) + kn, kr, kv_fp32 = split_rows(case["kv_fp32"]) + n, h, _ = case["q_fp32"].shape + return { + "kernel": { + "q_nope": qn.reshape(n, h, _FP8_D_NOPE_PADDED), + "q_rope": qr.reshape(n, h, _FP8_D_ROPE), + "unified_kv_nope": ukn, + "unified_kv_rope": ukr, + "kv_indices_prefix": case["kv_indices_prefix"], + "kv_indptr_prefix": case["kv_indptr_prefix"], + "kv_nope": kn, + "kv_rope": kr, + "kv_indices_extend": case["kv_indices_extend"], + "kv_indptr_extend": case["kv_indptr_extend"], + "attn_sink": case["attn_sink"], + }, + "ref": { + "q_fp32": q_fp32.reshape(n, h, _FP8_D_HEAD), + "ukv_fp32": ukv_fp32, + "kv_fp32": kv_fp32, + "kv_indices_prefix": case["kv_indices_prefix"], + "kv_indptr_prefix": case["kv_indptr_prefix"], + "kv_indices_extend": case["kv_indices_extend"], + "kv_indptr_extend": case["kv_indptr_extend"], + "attn_sink": case["attn_sink"], + }, + } - def _csr(total_rows: int, seed_offset: int): - if mode == "sparse": - return _random_csr( - n, - total_rows, - device=device, - kv_tile_size=_FP8_KV_TILE_SIZE, - seed=seed * 2 + seed_offset, - ) - if mode == "dense": - return _dense_csr(n, total_rows, device=device) - return _empty_csr(n, device=device) - ip_p, ix_p = _csr(total_pages, 1) - ip_e, ix_e = _csr(total_tokens, 2) +# --------------------------------------------------------------------------- +# Portable Triton candidate. +# --------------------------------------------------------------------------- - kernel = { - "q_nope": qn, - "q_rope": qr, - "unified_kv_nope": ukn, - "unified_kv_rope": ukr, - "kv_indices_prefix": ix_p, - "kv_indptr_prefix": ip_p, - "kv_nope": kn, - "kv_rope": kr, - "kv_indices_extend": ix_e, - "kv_indptr_extend": ip_e, - "attn_sink": attn_sink, - } - ref = { - "q_fp32": q_fp32, - "ukv_fp32": ukv_fp32, - "kv_fp32": kv_fp32, - "kv_indices_prefix": ix_p, - "kv_indptr_prefix": ip_p, - "kv_indices_extend": ix_e, - "kv_indptr_extend": ip_e, - "attn_sink": attn_sink, - } - return {"kernel": kernel, "ref": ref} +_ASM_HEADS = 128 +_ERR_TOL = 0.05 + + +def _merge_two_sources(ukv, kv, ix_p, ip_p, ix_e, ip_e): + """Merge the two CSR sources for the portable Triton candidate.""" + total_pages = ukv.shape[0] + pool = torch.cat([ukv, kv], dim=0) + lens_p = (ip_p[1:] - ip_p[:-1]).to(torch.int64) + lens_e = (ip_e[1:] - ip_e[:-1]).to(torch.int64) + indptr = torch.zeros(ip_p.numel(), dtype=torch.int32, device=ukv.device) + indptr[1:] = torch.cumsum(lens_p + lens_e, 0).to(torch.int32) + parts = [] + pp = ip_p.to(torch.int64).tolist() + pe = ip_e.to(torch.int64).tolist() + for i in range(len(pp) - 1): + parts.append(ix_p[pp[i] : pp[i + 1]]) + parts.append(ix_e[pe[i] : pe[i + 1]] + total_pages) + indices = ( + torch.cat(parts).to(torch.int32) + if parts + else torch.zeros(0, dtype=torch.int32, device=ukv.device) + ) + return pool, indices, indptr + + +def _run_triton(q, pool, indices, indptr, attn_sink, softmax_scale): + """Run the portable Triton sparse-prefill candidate.""" + out = torch.empty_like(q) + num_queries, num_heads, head_dim = q.shape + + def grid(META): + return (num_queries, triton.cdiv(num_heads, META["BLOCK_H"])) + + _sparse_attn_prefill_kernel[grid]( + q, + pool, + indices, + indptr, + attn_sink, + out, + q.stride(0), + q.stride(1), + q.stride(2), + pool.stride(0), + pool.stride(1), + out.stride(0), + out.stride(1), + out.stride(2), + num_heads, + head_dim, + pool.shape[0], + softmax_scale, + HAS_ATTN_SINK=True, + BLOCK_D=triton.next_power_of_2(head_dim), + ) + return out # --------------------------------------------------------------------------- @@ -528,14 +539,14 @@ def _get_tolerances(prec: str) -> tuple[float, float]: # --------------------------------------------------------------------------- -# Single-case driver -- both pytest and CLI go through this. +# Single-case driver -- all candidates share this one logical input case. # `@benchmark()` collects the kwargs into a row dict and merges in whatever # this function returns, so the CLI can build a pandas DataFrame. # --------------------------------------------------------------------------- @benchmark() -def run_pa_sparse_prefill_opus( +def run_sparse_prefill( n: int, h: int, d: int, @@ -544,6 +555,8 @@ def run_pa_sparse_prefill_opus( prec: str, *, mode: str = "sparse", + nnz_prefix: int | None = None, + nnz_extend: int | None = None, seed: int = 0, verify: bool = True, bench: bool = True, @@ -558,46 +571,119 @@ def run_pa_sparse_prefill_opus( f"prec={prec} mode={mode}]" ) + case = _make_sparse_case( + n, + h, + d, + total_pages, + total_tokens, + mode=mode, + nnz_prefix=nnz_prefix, + nnz_extend=nnz_extend, + seed=seed, + ) + nnz_p = int(case["kv_indices_prefix"].numel()) + nnz_e = int(case["kv_indices_extend"].numel()) + total_nnz = nnz_p + nnz_e + candidates = [] + if prec == "fp8": - data = _make_inputs_fp8(n, h, total_pages, total_tokens, mode=mode, seed=seed) - kernel_inputs = data["kernel"] - kernel_fn = pa_sparse_prefill_fp8_opus - ref_fn, ref_inputs = _ref_pa_sparse_prefill_fp8, data["ref"] + split = _make_split_inputs(case) + candidates.append( + ( + "opus", + lambda: pa_sparse_prefill_fp8_opus( + **split["kernel"], softmax_scale=softmax_scale + ), + _ref_pa_sparse_prefill_fp8(**split["ref"], softmax_scale=softmax_scale), + "fp8", + 1, + ) + ) else: - kernel_inputs = _make_inputs( - n, - h, - d, - total_pages, - total_tokens, - _PREC_TO_DTYPE[prec], - mode=mode, - seed=seed, + single = _make_single_inputs(case, _PREC_TO_DTYPE[prec]) + candidates.append( + ( + "opus", + lambda: pa_sparse_prefill_opus(**single, softmax_scale=softmax_scale), + _ref_pa_sparse_prefill_opus(**single, softmax_scale=softmax_scale), + prec, + torch.tensor([], dtype=_PREC_TO_DTYPE[prec]).element_size(), + ) ) - kernel_fn = pa_sparse_prefill_opus - ref_fn, ref_inputs = _ref_pa_sparse_prefill_opus, kernel_inputs - - nnz_p = int(kernel_inputs["kv_indices_prefix"].numel()) - nnz_e = int(kernel_inputs["kv_indices_extend"].numel()) - row: dict = {"nnz_prefix": nnz_p, "nnz_extend": nnz_e} - - if verify: - ref = ref_fn(**ref_inputs, softmax_scale=softmax_scale) - got = kernel_fn(**kernel_inputs, softmax_scale=softmax_scale) - rtol, atol = _get_tolerances(prec) - checkAllclose(got, ref, rtol=rtol, atol=atol, msg=msg) - - if bench: - # `@perftest()` returns (data, avg_us_per_iter). - _, lat_us = _profile_func( - kernel_fn, **kernel_inputs, softmax_scale=softmax_scale + + bf16_inputs = _make_single_inputs(case, torch.bfloat16) + pool, merged_indices, merged_indptr = _merge_two_sources( + bf16_inputs["unified_kv"], + bf16_inputs["kv"], + case["kv_indices_prefix"], + case["kv_indptr_prefix"], + case["kv_indices_extend"], + case["kv_indptr_extend"], + ) + triton_ref = _ref_pa_sparse_prefill_opus(**bf16_inputs, softmax_scale=softmax_scale) + candidates.append( + ( + "triton", + lambda: _run_triton( + bf16_inputs["q"], + pool, + merged_indices, + merged_indptr, + case["attn_sink"], + softmax_scale, + ), + triton_ref, + "bf16", + 2, + ) + ) + + if h == _ASM_HEADS and _get_gpu_arch() == "gfx1250": + split = _make_split_inputs(case) + candidates.append( + ( + "asm", + lambda: mla_sparse_prefill_fp8_asm( + **split["kernel"], softmax_scale=softmax_scale + ), + _ref_pa_sparse_prefill_fp8(**split["ref"], softmax_scale=softmax_scale), + "fp8", + 1, + ) ) - # Sparse attention FLOPS: 4 * H * total_nnz * D - total_nnz = nnz_p + nnz_e - flops = 4.0 * h * total_nnz * d - tflops = flops / max(lat_us * 1e-6, 1e-12) / 1e12 - row["latency_us"] = round(float(lat_us), 2) - row["TFLOPS"] = round(float(tflops), 2) + + # Report per-row nnz rather than the pool-wide total, so the column matches + # the --nnz-prefix/--nnz-extend the sweep was asked for instead of scaling + # with N. Exact for fixed/dense/empty (every row has the same count); a mean + # for sparse, whose rows vary. total_nnz below keeps the full count -- the + # FLOPs and bytes figures need the real work done, not the per-row figure. + row: dict = {"nnz_prefix": nnz_p // n, "nnz_extend": nnz_e // n} + for name, invoke, ref, candidate_prec, kv_esz in candidates: + if verify: + rtol, atol = _get_tolerances(candidate_prec) + err = checkAllclose( + invoke(), + ref, + rtol=rtol, + atol=atol, + tol_err_ratio=_ERR_TOL, + msg=f"{name}: {msg}", + ) + row[f"{name} err"] = err + + if bench: + _, lat_us = _profile_func(invoke) + flops = 4.0 * h * total_nnz * d + tflops = flops / max(lat_us * 1e-6, 1e-12) / 1e12 + row[f"{name} us"] = round(float(lat_us), 2) + row[f"{name} TFLOPS"] = round(float(tflops), 2) + row[f"{name} TB/s"] = round( + (n * h * d * kv_esz + total_nnz * d * kv_esz + n * h * d * 2) + / max(lat_us, 1e-12) + / 1e6, + 2, + ) return row @@ -625,9 +711,9 @@ def run_pa_sparse_prefill_opus( ids=lambda v: "x".join(map(str, v)) if isinstance(v, tuple) else str(v), ) @pytest.mark.parametrize("mode", _PYTEST_MODES) -def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): +def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): # bench=False keeps pytest fast; CLI path does the timing. - run_pa_sparse_prefill_opus( + run_sparse_prefill( n=n, h=h, d=512, @@ -649,8 +735,8 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): parser = argparse.ArgumentParser( formatter_class=argparse.RawTextHelpFormatter, description=( - "pa_sparse_prefill_opus correctness + benchmark driver.\n" - "All list arguments are swept via itertools.product." + "PA OPUS and gfx1250 asm MLA sparse-prefill correctness + benchmark " + "driver.\nAll list arguments are swept via itertools.product." ), ) parser.add_argument( @@ -658,15 +744,20 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): "--n_tokens", type=int, nargs="*", - default=[1024, 4096], - help="number of query tokens N (default: [1024, 4096])", + default=[512, 1024, 2048, 4096], + help="number of query tokens N (default: [512, 1024, 2048, 4096])", ) parser.add_argument( "--h_q", type=int, nargs="*", - default=[16, 32, 64, 128], - help="number of query heads H_Q (default: [16, 32, 64, 128])", + default=[128], + help=( + "number of query heads H_Q (default: [128]).\n" + "The asm candidate only registers at H_Q=128, so the default keeps the\n" + "three-way opus/triton/asm comparison. Pass 16/32/64 to sweep the\n" + "opus-vs-triton pair at other head counts." + ), ) parser.add_argument( "-d", @@ -679,10 +770,11 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): "--total_pages", type=int, nargs="*", - default=[4096, 16384], + default=[], help=( - "rows in unified_kv (default: [1024, 4096, 16384]). " - "Pass 0 to mirror -n for that sweep point." + "rows in unified_kv. Pass 0 to mirror -n for that sweep point.\n" + "Empty by default, which switches the mode/total_pages sweep off so a\n" + "bare run only does the explicit-nnz sweep; pass values to enable it." ), ) parser.add_argument( @@ -695,10 +787,11 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): "--prec", type=str, nargs="*", - default=["bf16", "fp8"], + default=["fp8"], choices=list(_PRECS), help=( - "precision(s) to sweep (default: [bf16, fp8]).\n" + "precision(s) to sweep (default: [fp8], the only one the asm\n" + "candidate implements).\n" " bf16/fp16: single-tensor Q/K/V/O kernel\n" " fp8 : split NoPE-fp8 / RoPE-bf16 DSA kernel" ), @@ -707,7 +800,7 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): "--mode", type=str, nargs="*", - default=["sparse", "dense"], + default=[], choices=list(_MODES), help=( "CSR mode(s) to sweep for both prefix and extend.\n" @@ -715,7 +808,7 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): " seeded at KV-tile boundaries (0, 1, T-1, T, T+1, ...)\n" " dense : every token sees every page / every kv row\n" " empty : all-empty CSR rows (sink-only output)\n" - "Default: [sparse, dense]." + "Empty by default (see --total_pages); this sweep needs both set." ), ) parser.add_argument( @@ -734,6 +827,43 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): default=0, help="RNG seed for input + CSR generation", ) +parser.add_argument( + "--nnz", + type=int, + nargs="*", + default=[], + help="explicit per-row nnz values for the shared boundary sweep", +) +parser.add_argument( + "--nnz-prefix", + type=int, + nargs="*", + default=[256, 1024, 4096, 8192, 16384], + help=( + "explicit per-row prefix nnz values for the shared shape sweep\n" + "(default: [256, 1024, 4096, 8192, 16384])" + ), +) +parser.add_argument( + "--nnz-extend", + type=int, + nargs="*", + default=[128], + help=( + "explicit per-row extend nnz values for the shared shape sweep\n" + "(default: [128])" + ), +) +parser.add_argument( + "--pool", + type=int, + default=4096, + help=( + "prefix/extend pool rows for explicit CSR sweeps. Raised to the largest\n" + "swept nnz when that is bigger, so the CSR never wraps and re-gathers\n" + "the same rows (default: 4096, i.e. 16384 for the default nnz sweep)." + ), +) if __name__ == "__main__": @@ -750,7 +880,7 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): # 0 is the sentinel for "mirror -n" on a per-sweep-point basis. total_pages = pages_arg if pages_arg > 0 else n total_tokens = args.total_tokens if args.total_tokens is not None else n - row = run_pa_sparse_prefill_opus( + row = run_sparse_prefill( n=n, h=h, d=args.head_dim, @@ -765,13 +895,59 @@ def test_pa_sparse_prefill_opus(prec, n, h, total_pages, total_tokens, mode): if row: rows.append(row) + if args.nnz: + pool = max(args.pool, max(args.nnz)) + for n, h, prec, nnz in itertools.product( + args.n_tokens, args.h_q, args.prec, args.nnz + ): + for npx, nex in ((nnz, 0), (0, nnz), (nnz, nnz)): + row = run_sparse_prefill( + n=n, + h=h, + d=args.head_dim, + total_pages=pool, + total_tokens=pool, + prec=prec, + mode="fixed", + nnz_prefix=npx, + nnz_extend=nex, + seed=args.seed, + verify=not args.no_verify, + bench=not args.no_bench, + ) + if row: + rows.append(row) + + if args.nnz_prefix and args.nnz_extend: + pool = max(args.pool, max(args.nnz_prefix), max(args.nnz_extend)) + for n, h, prec, npx, nex in itertools.product( + args.n_tokens, + args.h_q, + args.prec, + args.nnz_prefix, + args.nnz_extend, + ): + row = run_sparse_prefill( + n=n, + h=h, + d=args.head_dim, + total_pages=pool, + total_tokens=pool, + prec=prec, + mode="fixed", + nnz_prefix=npx, + nnz_extend=nex, + seed=args.seed, + verify=not args.no_verify, + bench=not args.no_bench, + ) + if row: + rows.append(row) + if rows: df = pd.DataFrame(rows) - # Drop columns that don't carry signal in the default sweep. drop_cols = [c for c in ("verify", "bench", "seed") if c in df.columns] if drop_cols: df = df.drop(columns=drop_cols) print() print(df.to_string(index=False)) - sys.exit(0) - sys.exit(0)