From e60084abcc8706eae6c65014cda7d4b828caa9bc Mon Sep 17 00:00:00 2001 From: junxiaguo Date: Wed, 26 Aug 2026 20:53:40 +0800 Subject: [PATCH 1/7] Add MLA v4 sparse prefill asm support Integrate the gfx1250 MLA implementation and consolidate sparse prefill correctness and performance tests. Co-authored-by: Cursor --- .github/scripts/split_tests.sh | 2 +- aiter/__init__.py | 1 + aiter/jit/optCompilerConfig.json | 11 + aiter/ops/mla_sparse_prefill.py | 143 +++++ csrc/py_itfs_cu/asm_mla_sparse_prefill.cu | 260 ++++++++ ...a_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co | Bin 0 -> 37936 bytes hsa/gfx1250/mla_v4/mla_v4_asm.csv | 11 + ...fill_opus.py => test_pa_sparse_prefill.py} | 570 ++++++++++++------ 8 files changed, 796 insertions(+), 202 deletions(-) create mode 100644 aiter/ops/mla_sparse_prefill.py create mode 100644 csrc/py_itfs_cu/asm_mla_sparse_prefill.cu create mode 100755 hsa/gfx1250/mla_v4/mla_a8w8_qh128_1tg_32mx4_32nx1_sparse_pfl.co rename op_tests/{test_pa_sparse_prefill_opus.py => test_pa_sparse_prefill.py} (59%) 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/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..bee74b2253 --- /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 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 0000000000000000000000000000000000000000..819856468874a2b287200a3c0c80901d4d42e194 GIT binary patch literal 37936 zcmeHw4_uYi{r|bob45f{M6w4(xse%}5#pb~^{tz;W`KrHFJw+YK~WJHpjltPbU`#$ z%rWO0YgTCgs9B@t8nwQFeO+VCHP+V|E7zQ>&26pFw5&Dn@BKdKIhX78Mg{Bp+g_7i z@7~XI&gXnSpL3q`oO{mad7k5K*JfU4I2>-LP=BU4EFAVR$!#CEG^m5?2Qe*Ggya9i zMWhHrGRwME_F3FvE7Kfq%Q0Eow3E)DjtF(CSwyLGmV*mI3{~f1?3wJ{q|P-5mZPuU zt#}u#g%5o>x1_YJ z@ZqG=6)P$VE5(CJ6{~J74F7J@`uyT`g}G~1l@vUtHPEA?6VH*BlvSA7b_410s+_DwL*ZwG}tbEn_{K`VV zAH9x>?}zd2U5z{_-`?p+)`HQO)lq-JHL^a7B438k23}dOE4MDay)A`3qKBp1ZcR zpwRJ{{fk$L1}?0824^dmms|!S@Kbo%Wh+-~n36I*S^qSyFD$QERa&x5|8A>a*R%i6 z{Ad$DuzH?ii(CO7iJ67hp69sjbjoXt!7;Nj6>%`FVC!qY-@9(=FcZgw{?kA8Jg1!K zS9|xb6s{624&YDD=0uUL!E4XELMiMSA;i@pvZHmR5aCX?&cOeMxVai0(~)`2F*l)z*9v$lqYd>G?YdYvaA(;_*_W|z&$HoI-kvN=z{SNgW4-SnB#mqETkk_Y<+ zOCI7IB6+B9sN`Y3VUkDqMo1p%8!35|Z*`SYNDULyY~Eh<6CD)0sYY zv!~Vru5mh(?;BCqK3oiX>olBaTbZ_8Ls*!S1J{Iw`Q^|&vafFP)wM6I_1qQhiS~~2 zj;a~$8SOpadw$JNM+vsWsUb1q1m>itGQd2)(>|{jJKDVB9^l?KpSZ6Ld0ww*!ud{M z7jQrD0PrC25b$tYwP*%E0z8T`-?hu^#`$BwsKhwyopcttaA`d=m@7x+AGt+C%om-O$-{<<~%65(S#oviKoSqvWU%v$P+8*81=JL6aclBv-J*wwW z=RPvdMc4(t3y?2Ff-eDOeLkEs-AM2z)H98F7y2%gc{oR&OI+x^u%2mlo|K>9n;>N- zG)<`YiGyvtfL`%(n+I)|?aP+36MYk<%*5=8^|D@DQLk;NSAB22oT%3zS+7U^^}5t| zDaxTmn3cD9O+(> z4J!eb8XzHj}s?v+$gW_?dbEq?5|SCP`i)&Mn^?Qq0Pe0=qjIhYZQ-{y!CcQt>~Fr z+jDG6__2}cLynC~UlJ9Sz9c$2-2=`B_m+=|Zh*dsW9Oy&Vq((SMoYlxO`Egv-+Sih zs`t#E+dQL3r)T#Z%W&g94>zN#hQMBg8B-EEN*9Z^ zwEGNa)#n~%Jwz<9s~BmfhDi9jMS378Bd0m(oL zFb$XiqyjU6D}k#3H!usB4O|Dz0n&j?U>=YKEC3b)9v}x;3@ibb0(rnPpa56_6alM& zVxR;l1ImF)U_Gz_xE0t4Q~_R~2B-zL0Cm6}Ks~S(XaKeW+kv})M&NE>2XHUY1ndGH z03HIGfk%Pez~ev*@FcJYcnW9*o&lZ(+JFw=Iba{~JkSZe0PF`|1iFBifCIompc{Av zI0U=~90ra6uLFMojskxM-UQwPdVu4=3E*Ae6z~D?A#nOO;oH*g|9uk<{H|dX@@3z} z38yC;ziAxD@$0SU0u9cb=von$;tLCl=5NR3!8z%koSbx>zgk>>z!-_&N&KduAh^2? zYDKm^wm603SI#Gz9p;d*bdGH-&oNDwAK?(Sm|tvyekaCT({zaHu!xk~BO;Q0;gQi^ z=&24Ho-%mI@ELW(hbJ3Do$1CfV1zSM)}cE4tMWG)24x+<+JBtOdU%|R)3cpRr0vVh z#Z^h>k}8gckHQDs3%DN)zzKu_;Xnis2@D6KfM_5F7!6p!7$6Q954eB?U;;1^NCYMU zlYt~48At)90W*M9U?y-Sa24PNW&yK->wq~xI*1G0bxz(T+SwZImj4!8rT2etwYz&2nza2L=B+zspi z?gg5FUBCmtLqId|D6kuN9B2Wa1oi+=0j;s+$I)N8}{lJSr7w{5r z05}MA1FrywfY*S-z!Bhe;19r2;LpIDz*|5Ma2z-RybGKHJ^(%hP6Gn>gaJ5#Fd!U= z03w0mKok%S!~mlK3m5~$0pkG|kN`{oCIX4TBw#X-1SA6~z%*b6kP6HMt^}?E+`ueg zHgFv<2S^7pfq6g{umD&Hcz_&WF|Y(!3giLHfC69zPz0<7ih&ZK3@8UGf%U)!;8tKG zPz88_8lV=~0@ML_0QJCDpaIwhYzOWF8iBik9l*Un6R-<-0C)&!1|9`=1CIkOz>~lp z;3=RLcm{YDXahQc=YV~{^FSx?0Q@{tnhXBUHN4K;K@e|-@?Ofw@IJ&!s3h@E?ln{UGROcU(asGE9 zK72J?h-8=IX>P@-d5W)O{tzMD%m>eAK6noE!I{D#%V)V1FLW!;$y2<9`H;(FKDdDS z;3DROi>V*6fimg`S5iNC1NDP9UW4&fTwwdCc$!;rYM$aNnGaoVm*UyX$N3!QgEK{f z?W5v_ZpAryikC1Sa(OPr1#cww>oo!@S-XA0zhQc)C2E=kwf&$*%Bx+fzH;GrfM^?K3v}#1GqqTZsA4 zF&{cCkL8Wcj;&cR|Cio*_u=5ZF7yekxv$yM{&8av<Xp7i=S_%3WDVQZ{(ruHJ{honKABkm*{pAUKCEr)?X|631M(c* z!}uiY+RtW9>hobeDt!uGm(nLMoa<{XI<|Vrol9z$)XKHzCQmKn3QKCf#B+cALC?<^ zBl94}kcs#~V7{F1T)_B*n#XyaltUx}>u~?tzivXkf1O;HT<;aHz_&NrR-(ULi2grc zyx3;p9*ni`K@0a_tbGs854s2E2i$||e#1(3uctb$L&mmtFTs7e^7ianPxfyR=DfE} z$E0)&WGG@Fk%-ag7?qA83_}cIB<_zel+kf5U1l)Kgrkg(bLlceP$rZQH?0h7_pMyi zi17r&2kv8zix}s}lp2gPJrl0i_EBnWUn`s`oQH-Fic04?KOF1saIDMEZ0GE+G9EOx zI%jjv?K!nMz8k;F>&x|i4&cF>|3)AiYx^7^8|(QSf%mknx89bq1KRRAozc{V_}Wb4 z?YA@D*D}Z78O2zE_UHXU$1{A#-^mcTzMMP1Zyvvmb7v3IIOb%77vla8j28@r-|%qr zjrdP2+|th2LJr~;HzKy+pL1iry$SI;W3V%wzR)J;;Ed~W4qj{L$@=R#HvE-yZ2HYL z0(_ElY_P|bp^$&HbzdG)aVdz{80TS>p0Do?%p`fm2Q?v_$TVKg)2N?@i1=mC69N^mu3Ij z&9#u2B*eJvao$U^FY*3k+hg{e-u9S1=bySg)cy7ux5v5uapwKP+=iA=S@s9KH_k?)Io|m8Cd>rQ|?RojToEKwW?&Z7~T+4YexQ_E;Fy;tyZrZ?n z%u%;9AKb`%@D4d|cGx-;*K+PHbttYEV8}MO6mMs~)S-9>=Zdt0d3Y1;fFGb8a5L?I zcQb#e5G~9H?_oZ;mHFUjwI5Ok^hzCyoAVU!rXNy=;yv_3>QMZw_TzmaI%o&Hk9NSF zv;*ER?Wp6Al(8@ucB4-3h}j3lAAP2XbVQ!8kBvw(@+6P{qSD{BkD1QvC0x_@zjG2r z#MJ%Xe8X6X-zp(wTDVT*ZxiKCyr<+Es#Ba=4rx-3X%1WNNM#i z%nh=w}P*A+B*JX>i|m~V5tKvb^O)V0hT(zQU{oG25~K~>Bd(lcVWgL zM3qDK8Jr(b=YrSViSzD&e4Zat=T7Fkddm-jU-H~8?;M2l3BBcIKF?X6dEq!djN?e~ zQClZ{NVL-@t8@qTQucLQ_Dv-_i1sPVbjl)q@dDAFkGYDRN7Z_23p_7ZH+Z_M3v#-v zX|q;Dq|`-3L=Q5ZRfEj1DnnfKa<rE@LzK!USUr`JxHN?c|Fm8s0GQC5c8g`O-;74f(Rk#wPORNk%PsO0uz; zJT=AGLY_X&xSf2(41;q|?c*XhW4HgDe)7J)32TuB9%s7;>-)g9zFr3yV$#9Lu``_} zc*F(QIi29>ix)TtfzP{aDO!KhL2GJ+%pY^baCvSHiEtwJJlViCGQWZAG}t-$xOwC@ zK3_XFC~UI%`yt`0Ma%)G&mpYxC)~#dziu9T!wmNw`Q}DmKc8Vl@p+sh!$7}w{L1mh zBM#r|!(>{RIlPK<`F}>g*q)KIyi2te@6Hox?K5$TY2!X~Q$(WSZ8`v=JHOWtz6dw5SZ1Ow)FmHZmhYrfDCT7M(Fc zrfGkeHY#JHOw+zGEhZxo_uVD?BgYxmc^Q+coH&lU$goCdOqS`!C5Cl=Mv_c-US?QU zMsihbMoQI~jA>P4GiFGAtRr<%k1@rtcs*H1rcs}Bx=Ldm{c=N88tW+220O!8j`fr0 z=xaP@-FVJ^$8**zP@Z+-d5;)%u*ZyWQf8oa()G~3%X9dPejgTa4}@u(w6E_kz4pnr zpM!Hf%(ZYf)}Va8%{421JS(5;TY|)+OKymcy1nuxg_R!axSUn zdo3=-XU_xe_+=E@lH>P4=5cZMJkAL{LdZE}hhEd_`Ga2T>-D{!Pb}o`r9Joa>%lz5 zrS$A$Y(URBbc{gHJ%s!XSMx8HZKy8Cc)bPnQS({G5cK?n>tH=c(Kb0p^Ou9oJXLN# zV-9*Aqhk+xJ|lnMF^}Y2PWrpE+<7g3+wFRDttIP>d5w-&$hzWuk(|S*IiJ6*yf=Dd z8M@xVb0i1L+VfZG^Im0F*EM(!rEPKE=$D5rUB}=#7U#YGc^2lNoS#py=e{Dj=he$_ z1B(>Nb&+E2%aLOZF$_FB-VkzLj=3z#gzGXrpxsV^Ds9)Zp_>`LHe^T zPR_COdTlSl{W(F}-WE5Ac9>_|abSMSa-GNQas|D%bUEh9a?FFRd+@%XA@0X-KzCcr zNc&j{^ziqBegx)Q(RM!UvTyt21@`-0iIe{|Vuf5ItdFvvH=&*m@yZY7I~0?!e(>UX z)+Rg)@FJd_lier~BRT7L^ArBA`gQINp>yfdborYM=5Mf$!|*qlzr{`Ps{_~48*9Wb zYEpC6xX<0{q%7yH@;FM40UaV*awq54@z-pT^Xr5!Y{SC==y2OQPGCM`=Tisg-WSci z*C}FT*)B0!ayNA+eo@WECuI+pIy&J;U2H=9BIw}Uow~IjoTm%;f$?D<{Q&Q$@13wK z%fNTPE&7f*J#9kn0BwQ~QN9Z{8e(Vp>l)~Tu3>3mVU(9bs4Yf zLb-0}ZH!$eeV4YOOZp03(pSnrmz068oNEc_k^Vss{o4td!_ec6t%+Z>Q_Txy{V6*K z_t9Yl#!!dD&?99|LJ!s}l;IpwWI_h(Tgn`Tp1Rmw@rxkC^$gPY|ApZ4qSs#R!orjEn=eNR?5wR?#|eR zaqL4}r%~<<8_WY1C9Qz;JZxHfi|E2Hv-6(@jnGaud+p+w(jG1!Hq}Sl=w;Q0VJNC#p z0a>=2)WyCIUw1;+24UyJmToiZ8lQ2c)D_$|x*vO?w<*>bFPgTr2e+5NJOT5r4q-^{ z6i&%qy4`d;u+Mr?7pVt*MYjvvL4?UN_!XDecOP}(I`3y4S*{NDn`0Bk3zU<-!UpmW zupIbMe8yE(?6Y;yPdVOiOcUX@ZxOa{k;*q$hP$c>ZQlUBEwQtt`~?}as>E>H50r!6 zeFEj^M?%KzDvo^(&@biT$AuZ!RSD>mexMxvn2<3?$~8i7YwR-VFKvZd|Ir?OZiG%L z3!iC6KOdle_NgZ5YKpBH&;CLE`nkbH8R=EL=1nLs^Ks25W@O5Iyyx@o#ToNtelyB9 z$L<=>K9QJ_g;<>Zj6iYiB^e8(oo48fa~qo&uBGGt z!?3uvc3o>&oUcQ+p1S#bfqu0>pOlAPx%R$Z$vfv7mP^TPr(7%aw#FWj{znQT)Kk7|24mls&21m~y;MkZb1k;JTi2iO`9$Fpgs=+gs}!8u09B zC*`0|%E30#iDF$c<}LKnu# zxOI@D9_f?fJv-MSU#R?i_{8$k7sb!Y7*YO~_J>FAlxrO~uJ=IV68c%vK;sgAJqtg5 zT!Ll!oa)oYCH&=Jlh>Zlh}oZT(jr6ln?v*ic3hJ@eGZ2 z*>B-*e4zMZ6{O`Qj&u zOUO3Zi)WiCZ`%kNm(b56uf=^E+&;n2?-oKvKeJ?9;;eR&{SWi2fVhNyt|^>I4~R=p z51(<$xCHYKsknrG)*1ZlkMRZl>{G@UT)lP}m(b5bWw{Bxwsbi@8@W&tB9J0)?+|6v+G*u|jfhTp}&uqvI0LjkbP( zy1~2IC+6U~bg~Sti(XH&-wEi4%pU3oKTFvx)V(Wh)|qjMBHQ29w!bKw1)Wk3{>r#S zvF&S#?JHz*psPD=S?IU~bV@&=Q~F7H=#=vCQ^qBrQ~C#;^pE`~4>~*2YGhm@@OiO6 zF2Q~j7?+^zNjy77oS?5~)Bdpa3HIN>xCHazdth8bKQlfnE|C@#mtbEHj7u;d{^>U3I!Zrl4sM$>;u2|w zj7#X}>;7lzZh4(M#4O33Vz%Ti-EO)aIPQ6IU7M%}eMPs6d>($CEYmHZ3)goab%FP@ zjx1M)Hfl~wkmaPWuz~yo0_DJm_*|TQwhsC!$Nnv!mCvz#OSgT?RKD@Kc@x^c0eV}~ zR9r$nJD+F!fpXFhl%pSfj?S^K0s5sp{Fl$tp-=jOa`c1G(ka&ny{&1>q`$NkDlS2L z^jXFwC<~uyN5>_o-yfGqtC8^l>eq1zzJG8M<(uGB6Z6qO_#T3O&ynvTFuxhTG^g#7 z`FtNi2z$M(xR&oH&`vY-NI8y^cwd2YcE%<6-U9PmprSW3!9Trrrd{Np8J9>qGcLh) z;x!SdUkCGW?Q~p%?Zv!s)D7d}8F2}=ANxo+%E`PY`5Po-9c)MDMWGzVzO&*IY*T@C zY+zi1ZAx2F(8n?y-`TcqC3i+#BCu^)#)2J;hrMx$fVQL^w87bN3Ch!F#w7yVRLP$m zmq@Gmc7-1OIrPLdGT1c73wA1m$pDKQ=Cr_T0zEB`5>mJ~l3qcBJpP1oHce zOFa4BuJ+-2evkVDq1pD+h`$^k96tH@kce;}z6s`gT!ThLL=DP_s4|B{hKG%aj0($$ ztm+v$JZkW;;bDUpA$C0C@JKhx`xq;8Wk>e2ZaSMgA|wGr`q{HH*C2ux69LVp!LaziL=>$X_?C zbn-V1E0g?f!C66<$2J!`_wT*m{X>BK8Vp?~RFEg!1@)Xm$n>^jbLmP0KY28cy57TNQUt?Ok z$k&?I1LW&X>ml-7(`qKqH?2p>i%e@b`9{-voctxzY9Z&F)|2Ey)7nE`Wm->>*P2!< z`OBvD40)YtJxjjDwA#r3Wm+BNYSVg-yxFw&k-uVE&y&AuTAk#to7M~DZ<^MA^0!Uv zMe;qS)kWTES}&0wG_3>VM@;J=`7zV#CO=_XuaJLWT8GF_o7QXO|2D0|TA?{+($ZC%<7@C&+&?t#`?PHLX+R zW2W^1`J`!mNPge6PLuy`TG!Y0jt4)#o?{cz9p}b_S!{=MD+^_?kv908f;RY^qCYk|LfCNO z83t|e83t|e`2=n7`9yzgU{9a)EWC;ld`vUBBsDmz!*)_*(R5O&;~l$}|(D?79ARCccWTL0}d zgs?N`8_G`lx0IdC?l2Z>KASoudC$c2>Wv>=eJM?3Dbr|8@?C zuv7MXWvBd)%1-59l%4f&_uo!W2s<0zQFd;9Pubb{H)W^F*MB?Ms_U$A$7`B$+^HFC z%5kT5s42&{EhDh=Z9jf*IKD-Nuv0fm*}3Cb5|8`s<>@-}c>}^z#S?Cict*?By-|90|1*lAg=>^!+r+1ay3*?Fq8|8~kk*lDd$cAmLe*?IOq zm7TWR`ftY@!cNB~W#_rum7RTeDm%}At^al!LfGm2hO+a*x0Idz-%)m6ys!UunnKv= zx?kCO>0xE(z;~6MgWv1Fot6-Gy1%dNyz)b3=g^Ooo!8p?Z>J-Kox}g5>>T-tvh(`S zl$}5PqW^ZfLfAR_zsk;^UsiVBd{x$m;4b2x;Zp5H4w$N#A8ocN2f^X}XIx6>2C z&Z&2loe$nqc0T-@vUA!O+74r{fol=QU>iPW&6A6G>&LBm7UiD3=E*|2Gi+srwiUSc z(YE?p`wYZZ;2KEV>T3-&5LuYT_PuW7GP_&0i-pYGr2iE>T|{x^E?Jgu*PqetMIJ>k}SityjD`5E2^xI@0Lgipl4 zUcml$9)k8AJGbW;_B8tk?v2L&_fdO}wY+103-XV+7HO$!|F_T0RQtdE$CYaTx2vyG z`@db|R{OtwewNz*?b_LD|Fsnm=cJH-t;JrG>x%cYMy;pbcy}Ez<_v&)6x6aVx zyV`TA!`grUS8z>O`%=fQ_FoukyjX!@_W`oYXhx_mL z9xlg=!N=5i@!L1mc(EA&Ud+!N_Wa^Ez3=65o-X&?6Y?7*uf5MpoxRUXgS~f4qrG=a zt-W_kJ@fJH(RO=}mK_KaR$YTRn7t>9mwT&V-d1;g7d#~+Rn!~9K4TT@E(1^`{o7ji^p%+2dX2Ad#7BB>w<52yx2SC zL=%qjEfA0S{ZSqEMQOq}LHyo`$NaXa0sEIwj^7ONnBM_4Vt*3Kaqkfx??gWL74c$U zk(2m_#)WhqCnBv5`-1E|srLThaWc{xu(t>0k|D?ARHR|=Eh)$SIC$(re(;`I!Fyvd zAJ@#k?-lcL?%&&r*FN6`J&E#KH;GA-n}heQ3f`A0cn>PsZu$7e2!3K;nQ?3nH+<*u z9He2-m~m{6IY{Gi7SgcqiD?%EN%aDfsPNZBJ(s;Z9Y1rFD%IW={c>l__OSbQu6v>_T-j`kW-k06>zLxv! zeJ%GhA8mKQ-mmgd9@`D$0QREna4GJz_gd)6Q|I0Jd48FFw# zWO>|-v}Wvq#I|cj8jo9$)`IJ|7FRgUB88`cJ)zEy1$>RZ)X=A(SQP~W6(XFkr= zI3BdGLGT_0-1h*o{{04m_Y?4s<3aldh@kxexUYcJWA6pPe4HN;_ILlmKt8UKf1iNhy#V~< zc#*AVwd7_|Z0jkJ+$xYS^&lU-oB7~90{P%)^;qPFpZE^9cO1v@9HikJ+um^;$5}|j zce1_XIF56W#(e?UZgL#Qv40%L_x)^-JU8{*<2b&>XL}SOtrN#(QfI(8j_>L>NSy)W zIKGefkK<*sZXNcXLY+eG71U+#71V9-1GLZH2WUU@QP%_Z{yv}0IF9e<*_XViQxlHs z;D0l|d1sr{A&tikNNd4&?UZXk8jl;1#&6OoC&zId`^Ry7KhCymf-Ju!Z-y+t1!vne zBaO!`NaOe4lxsm6k6V$}if^(hC&zJqpB*@kW1kfI&;fZ|b2W}*-xK=Ki8NeuHI8GS z63TTU4cAa@pk?A1Xz zIgaC4kK@CSaSXce9R3O0APfT#|K{K$cm83IIQeJUJjmt)c0RsCmw#n8pFHg=am8&X?)-6j~ zUb^=31vh+t?d8SA>(^dZR$jWgaCv1CGF?}!NSl(JpPIIOS@M*Ea5W%5srZURsb}nJHM{; zmLs{axV$vCEWe;&Rmn=hYOgIUfdF(><}U*mp|mu{fY^ov-`3W|u9f^R@q)F(~5S6zVOn?<>vvzOvQGI_P8a#Q^K2 z_3Qqq*^RuwKX>2>j%k#nmtI 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,183 @@ 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 + 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 ) - attn_sink = torch.randn(h, device=device, dtype=torch.float32) * 0.25 + 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, seed_offset: int): + 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 +541,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 +557,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 +573,126 @@ 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) + 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 bench: - # `@perftest()` returns (data, avg_us_per_iter). - _, lat_us = _profile_func( - kernel_fn, **kernel_inputs, softmax_scale=softmax_scale + 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 + ), + split["ref"], + "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) + + row: dict = {"nnz_prefix": nnz_p, "nnz_extend": nnz_e} + 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 +720,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 +744,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( @@ -734,6 +829,33 @@ 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=[], + help="explicit per-row prefix nnz values for the shared shape sweep", +) +parser.add_argument( + "--nnz-extend", + type=int, + nargs="*", + default=[], + help="explicit per-row extend nnz values for the shared shape sweep", +) +parser.add_argument( + "--pool", + type=int, + default=4096, + help="prefix/extend pool rows for explicit CSR sweeps", +) if __name__ == "__main__": @@ -750,7 +872,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 +887,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) From 124681de5a39431f500a8c522236f5b83403c434 Mon Sep 17 00:00:00 2001 From: junxiaguo Date: Wed, 26 Aug 2026 20:58:21 +0800 Subject: [PATCH 2/7] Update op_tests/test_pa_sparse_prefill.py Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> --- op_tests/test_pa_sparse_prefill.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/op_tests/test_pa_sparse_prefill.py b/op_tests/test_pa_sparse_prefill.py index 4794a0737e..d2c0bd4b97 100644 --- a/op_tests/test_pa_sparse_prefill.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -30,7 +30,7 @@ from aiter.ops.triton._triton_kernels.attention.sparse_attention_dsv4 import ( _sparse_attn_prefill_kernel, ) -from aiter.test_common import benchmark, checkAllclose, perftest, run_perftest +from aiter.test_common import benchmark, checkAllclose, perftest # --------------------------------------------------------------------------- # Skip helpers From 3abbd1ae4454398595198939aed736a191367509 Mon Sep 17 00:00:00 2001 From: junxiaguo Date: Wed, 26 Aug 2026 13:19:57 +0000 Subject: [PATCH 3/7] Apply black formatting to test_pa_sparse_prefill Pure reformat, no behaviour change. Fixes the failing black CI job (black[colorama]==26.5.1, default line length). --- op_tests/test_pa_sparse_prefill.py | 28 +++++++--------------------- 1 file changed, 7 insertions(+), 21 deletions(-) diff --git a/op_tests/test_pa_sparse_prefill.py b/op_tests/test_pa_sparse_prefill.py index d2c0bd4b97..b3fdac1f9d 100644 --- a/op_tests/test_pa_sparse_prefill.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -362,9 +362,7 @@ def _make_sparse_case( 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 - ) + 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 @@ -414,7 +412,7 @@ 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) + rope = flat[:, _FP8_D_NOPE:].to(torch.bfloat16) return nope, rope, torch.cat([deq, rope.to(torch.float32)], dim=1) qn, qr, q_fp32 = split_rows(case["q_fp32"]) @@ -597,9 +595,7 @@ def run_sparse_prefill( lambda: pa_sparse_prefill_fp8_opus( **split["kernel"], softmax_scale=softmax_scale ), - _ref_pa_sparse_prefill_fp8( - **split["ref"], softmax_scale=softmax_scale - ), + _ref_pa_sparse_prefill_fp8(**split["ref"], softmax_scale=softmax_scale), "fp8", 1, ) @@ -609,12 +605,8 @@ def run_sparse_prefill( candidates.append( ( "opus", - lambda: pa_sparse_prefill_opus( - **single, softmax_scale=softmax_scale - ), - _ref_pa_sparse_prefill_opus( - **single, softmax_scale=softmax_scale - ), + 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(), ) @@ -629,9 +621,7 @@ def run_sparse_prefill( case["kv_indices_extend"], case["kv_indptr_extend"], ) - triton_ref = _ref_pa_sparse_prefill_opus( - **bf16_inputs, softmax_scale=softmax_scale - ) + triton_ref = _ref_pa_sparse_prefill_opus(**bf16_inputs, softmax_scale=softmax_scale) candidates.append( ( "triton", @@ -684,11 +674,7 @@ def run_sparse_prefill( 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 - ) + (n * h * d * kv_esz + total_nnz * d * kv_esz + n * h * d * 2) / max(lat_us, 1e-12) / 1e6, 2, From 432a1d0d99ab3d161a488be930b776cbdc69d984 Mon Sep 17 00:00:00 2001 From: junxiaguo Date: Wed, 26 Aug 2026 13:20:11 +0000 Subject: [PATCH 4/7] Fix asm candidate reference in test_pa_sparse_prefill The asm candidate passed split["ref"] -- the raw input dict -- where checkAllclose expects the reference tensor, so the first asm comparison died with: TypeError: isclose(): argument 'other' (position 2) must be Tensor, not dict meaning the asm path could never run. Compute the fp8 reference the same way the opus fp8 candidate above it does. --- op_tests/test_pa_sparse_prefill.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/op_tests/test_pa_sparse_prefill.py b/op_tests/test_pa_sparse_prefill.py index b3fdac1f9d..35b79a2f7c 100644 --- a/op_tests/test_pa_sparse_prefill.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -647,7 +647,7 @@ def run_sparse_prefill( lambda: mla_sparse_prefill_fp8_asm( **split["kernel"], softmax_scale=softmax_scale ), - split["ref"], + _ref_pa_sparse_prefill_fp8(**split["ref"], softmax_scale=softmax_scale), "fp8", 1, ) From 820c3f407d36bcee4af25d80c75d68d19728c382 Mon Sep 17 00:00:00 2001 From: junxiaguo Date: Wed, 26 Aug 2026 13:20:21 +0000 Subject: [PATCH 5/7] Report per-row nnz and default the CLI to the asm comparison sweep Two test-driver changes: * nnz_prefix/nnz_extend columns now report per-row nnz instead of the pool-wide total, so they match the --nnz-prefix/--nnz-extend asked for rather than scaling with N. total_nnz still carries the full count -- the TFLOPS/TB-s figures need the real work done. * CLI defaults now describe the three-way opus/triton/asm comparison out of the box: N in [512, 1024, 2048, 4096] x nnz_prefix in [256, 1024, 4096, 8192, 16384] x nnz_extend 128, at H_Q=128 fp8 (the only shape the asm candidate registers for). --mode/--total_pages default empty so the unrelated mode sweep stays off unless asked for. Every flag still overrides. Pytest coverage is unaffected: it reads _PYTEST_SHAPES/_PYTEST_MODES, not argparse. --- op_tests/test_pa_sparse_prefill.py | 56 +++++++++++++++++++++--------- 1 file changed, 39 insertions(+), 17 deletions(-) diff --git a/op_tests/test_pa_sparse_prefill.py b/op_tests/test_pa_sparse_prefill.py index 35b79a2f7c..866638142d 100644 --- a/op_tests/test_pa_sparse_prefill.py +++ b/op_tests/test_pa_sparse_prefill.py @@ -653,7 +653,12 @@ def run_sparse_prefill( ) ) - row: dict = {"nnz_prefix": nnz_p, "nnz_extend": nnz_e} + # 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) @@ -739,15 +744,20 @@ def test_pa_sparse_prefill(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", @@ -760,10 +770,11 @@ def test_pa_sparse_prefill(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( @@ -776,10 +787,11 @@ def test_pa_sparse_prefill(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" ), @@ -788,7 +800,7 @@ def test_pa_sparse_prefill(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" @@ -796,7 +808,7 @@ def test_pa_sparse_prefill(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( @@ -826,21 +838,31 @@ def test_pa_sparse_prefill(prec, n, h, total_pages, total_tokens, mode): "--nnz-prefix", type=int, nargs="*", - default=[], - help="explicit per-row prefix nnz values for the shared shape sweep", + 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=[], - help="explicit per-row extend nnz values for the shared shape sweep", + 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", + 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)." + ), ) From 617fc3441eb3ac17866c46060f736f2673d1e553 Mon Sep 17 00:00:00 2001 From: junxiaguo Date: Wed, 26 Aug 2026 14:02:55 +0000 Subject: [PATCH 6/7] Accept an over-allocated CSR indptr in mla_sparse_prefill check_csr required indptr->numel() == T+1 exactly. Decode reuses this kernel with the extend region empty and sizes its CSR row-pointer buffers once at [max_batch+1], launching with the live batch, so numel > T+1 is the normal case there rather than a mistake -- and the exact test rejected it outright. The kernel reads indptr[0..T] and nothing past it, so the extra tail is inert: verified bit-identical output against the exactly-sized call. An undersized indptr is still rejected. Trade-off: an indptr built for a different T is no longer caught here. Separating that from the legitimate case needs device data (indptr[T] against the indices length), i.e. a sync per call. Callers that can slice to [:T+1] should. --- csrc/py_itfs_cu/asm_mla_sparse_prefill.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/csrc/py_itfs_cu/asm_mla_sparse_prefill.cu b/csrc/py_itfs_cu/asm_mla_sparse_prefill.cu index bee74b2253..790f916ef0 100644 --- a/csrc/py_itfs_cu/asm_mla_sparse_prefill.cu +++ b/csrc/py_itfs_cu/asm_mla_sparse_prefill.cu @@ -89,9 +89,9 @@ void check_csr(const aiter_tensor_t* indptr, { 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 T+1 = ", t + 1, - " elements, got ", indptr->numel()); + 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 From a5cc34e1d1092f6965fbdf007fa882d339e73f2d Mon Sep 17 00:00:00 2001 From: junxiaguo Date: Thu, 27 Aug 2026 10:46:16 +0000 Subject: [PATCH 7/7] Fix int32 overflow in sparse prefill query offset `_sparse_attn_prefill_kernel` derived `query_idx` from `tl.program_id(0)`, which Triton types as int32. The q/out addresses are computed as `query_idx * q_stride_t` and `query_idx * out_stride_t`, and in the V4 layout that stride is `num_heads * head_dim` = 128 * 512 = 65536. The product therefore leaves the int32 positive range at `query_idx >= 32768` and wraps to a negative offset, so the kernel reads and writes outside the q/out allocations. Observed as NaNs followed by a hard GPU page fault: Memory access fault by GPU node-2 ... Reason: Page not present Verified on gfx1250 with a fixed-pattern sparse prefill case (H=128, D=512, pool=16384, nnz_prefix=256, nnz_extend=128): N=32768 before: clean (largest size that still fits int32) N=32769 before: fault after: nan=0 inf=0 N=65536 before: fault after: nan=0 inf=0 Promoting `query_idx` to int64 moves both offsets to 64-bit address arithmetic. This mirrors the existing `slot_off` cast a few lines below, which already handles the same class of overflow on the pool index; the difference is that the wrapped pool offset stays inside the allocation and reads silently, while this one faults. --- .../_triton_kernels/attention/sparse_attention_dsv4.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) 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)