diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py index 35e8c1be49..7d653e572e 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py @@ -407,6 +407,14 @@ class FmhaFwdPipeline: n += "_nqscale" n += "_" + self.F_kv_memory_layout + "_" + self.F_kv_lookup_table + # StreamLLM sink token. These batch-prefill pipelines have no sink + # dimension (always no-sink), but the aiter JIT filter + # (cmdGenFunc_mha_batch_prefill, PR #2794) appends `_nsink*`/`_sink*` + # to the kernel filter. Without a sink token in the generated name the + # fnmatch matches zero kernels → empty dispatcher. The C++ inner + # dispatch never keys on has_sink, so emitting `_nsink` here is purely + # to satisfy the filter and is semantically correct (no-sink). + n += "_nsink" return n diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_ps128.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_ps128.cpp new file mode 100644 index 0000000000..72127f2a8c --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_ps128.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS(decode_d128_m128, bf16, true, 128) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_s_ps128.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_s_ps128.cpp new file mode 100644 index 0000000000..f0a81e6869 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_s_ps128.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS(decode_d128_m32, bf16, true, 128) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_t_ps128.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_t_ps128.cpp new file mode 100644 index 0000000000..0f08fae5af --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_mask_decode_t_ps128.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS(decode_d128_m16, bf16, true, 128) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_ps128.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_ps128.cpp new file mode 100644 index 0000000000..1802fd02bb --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_ps128.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS(decode_d128_m128, bf16, false, 128) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_s_ps128.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_s_ps128.cpp new file mode 100644 index 0000000000..a70714e608 --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_s_ps128.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS(decode_d128_m32, bf16, false, 128) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_t_ps128.cpp b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_t_ps128.cpp new file mode 100644 index 0000000000..b4848856ca --- /dev/null +++ b/example/ck_tile/42_unified_attention/instances/unified_attention_d128_bf16_nmask_decode_t_ps128.cpp @@ -0,0 +1,11 @@ +// SPDX-License-Identifier: MIT +// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "unified_attention.hpp" +#include "unified_attention_impl.hpp" + +namespace ck_tile { + +INST_UNIFIED_ATTENTION_DISPATCH_PS(decode_d128_m16, bf16, false, 128) + +} // namespace ck_tile diff --git a/example/ck_tile/42_unified_attention/unified_attention.cpp b/example/ck_tile/42_unified_attention/unified_attention.cpp index d9f14b3e73..4ff54f4704 100644 --- a/example/ck_tile/42_unified_attention/unified_attention.cpp +++ b/example/ck_tile/42_unified_attention/unified_attention.cpp @@ -68,6 +68,21 @@ static bool ua_prefill_fallback_enabled() return enabled; } +// Diagnostic A/B knob: force every dispatch onto the PageSize=0 runtime +// (non-strength-reduced) address path even when a compile-time-pinned +// page-size instance exists. Lets us measure the address-op win of the +// ps16/32/64/128 menu (esp. the new ps128 decode instances) on identical +// shapes without recompiling. OFF by default (uses the pinned instances). +// AITER_UA_FORCE_RUNTIME_PS=1 -> always dispatch PageSize=0 +static bool ua_force_runtime_ps_enabled() +{ + static const bool enabled = [] { + const char* e = std::getenv("AITER_UA_FORCE_RUNTIME_PS"); + return e != nullptr && e[0] != '\0' && e[0] != '0'; + }(); + return enabled; +} + KernelConfig select_config(const unified_attention_args& args) { KernelConfig cfg; @@ -221,6 +236,9 @@ template dispatch_page_size(const unified_attention_args& args, const stream_config& config) { + if(ua_force_runtime_ps_enabled()) + return dispatch_one(args, config); + if constexpr(V == KernelVariant::prefill_d128 || V == KernelVariant::prefill_d64) { @@ -232,11 +250,31 @@ std::pair dispatch_page_size(const unified_attention_args& args, default: return dispatch_one(args, config); } } + else if constexpr((V == KernelVariant::decode_d128_m128 || + V == KernelVariant::decode_d128_m32 || + V == KernelVariant::decode_d128_m16) && + DType == unified_attention_args::data_type_enum::bf16) + { + // page_size=128 decode menu (d128, bf16). Even though the decode tiers + // skip the 8-warp Tier-0/Tier-2 scalar-promote gate, pinning page_size + // at compile time still collapses every `/ page_size`, `% page_size`, + // `* page_size` in the per-tile address chain to a shift/mask (the + // div-by-128 -> `shr 7`), which the profile flagged as the dominant + // SALU cost. Compiled only for ps=128 here so we can A/B the decode + // address-op win and compare to the legacy split-KV kernel (which + // requires page_size % 128 == 0) on equal footing. Other page sizes + // and dtypes fall through to the runtime-page-size catch-all. + switch(args.page_blk_size) + { + case 128: return dispatch_one(args, config); + default: return dispatch_one(args, config); + } + } else { - // Decode variants stay on the runtime-page-size path — they don't - // benefit from the Tier-0/Tier-2 gate (too few warps for the scalar- - // promote payoff) and the binary-size cost isn't justified. + // Remaining decode variants stay on the runtime-page-size path — they + // don't benefit from the Tier-0/Tier-2 gate (too few warps for the + // scalar-promote payoff) and the binary-size cost isn't justified yet. return dispatch_one(args, config); } } diff --git a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp index 9bf931262c..71f2dbb1e2 100644 --- a/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp +++ b/include/ck_tile/ops/unified_attention/kernel/unified_attention_kernel.hpp @@ -447,8 +447,33 @@ struct UnifiedAttentionKernel index_t num_blocks = total_num_kv_blocks; if(kargs.num_splits > 1) { - const index_t blocks_per_split = - ck_tile::max(index_t(1), (total_num_kv_blocks + kargs.num_splits - 1) / kargs.num_splits); + // The split PARTITION (blocks_per_split + start) must be identical + // across every query tile of this sequence. With a GQA pack where + // kBlockM % num_queries_per_kv != 0 (e.g. d=128, qpkv=6) the tile's + // last 1-2 MFMA rows spill into the *next* query tile's first token, + // so that token is co-owned: query tile N (spill row) and tile N+1 + // (real row) both store the same (token, split) workspace slot. + // + // Deriving blocks_per_split from the *per-tile causal* horizon + // (total_num_kv_blocks, which grows with q_block_local_idx) makes + // split s cover a different KV-block range in tile N vs tile N+1. + // The two co-owned stores then hold partials computed over different + // ranges and race non-deterministically -> a ~1-row error on the + // tile-boundary token (observed only under split-KV + causal + + // non-dividing GQA ratio; MHA and ratios that divide kBlockM are + // immune because no token is shared across tiles). + // + // Fix: partition over the causal-INDEPENDENT full sequence block + // count so split s maps to the same blocks in every tile, then clamp + // only the END by the per-tile causal horizon. The extra blocks an + // earlier tile would skip are fully masked per-pixel for the shared + // token, so both co-owned stores compute the identical partial and + // the duplicate store is idempotent again. For num_splits == 1 this + // path is not taken, so the non-split behaviour is unchanged. + const index_t full_num_kv_blocks = + amd_wave_read_first_lane((seq_len + kPageBlockSize - 1) / kPageBlockSize); + const index_t blocks_per_split = ck_tile::max( + index_t(1), (full_num_kv_blocks + kargs.num_splits - 1) / kargs.num_splits); num_blocks_start = ck_tile::min(blocks_per_split * i_split, total_num_kv_blocks); num_blocks = ck_tile::min(blocks_per_split * (i_split + 1), total_num_kv_blocks); if(num_blocks_start >= num_blocks)