CK-UA: fix split-KV partition for non-dividing GQA + add ps128 decode instances

Under split-KV, a KV token co-owned by two query tiles (which happens only
when num_queries_per_kv does not divide kBlockM, e.g. d=128 qpkv=6) was
assigned its split partition from the per-tile causal horizon
(total_num_kv_blocks, which grows with the query tile index). The two owning
tiles then reduced disjoint KV-block ranges for that shared token and the
combine step merged partials over different ranges -> a ~1-row error / NaN on
the tile-boundary token. MHA and ratios that divide kBlockM are immune (no
token is shared across tiles).

Fix: derive blocks_per_split from the causal-INDEPENDENT full-sequence block
count so split s maps to the same blocks in every query tile, then clamp only
the END by the per-tile causal horizon. The duplicate co-owned store becomes
idempotent again. num_splits == 1 is unchanged.

Also adds the d128 bf16 page_size=128 decode instances (mask/nmask x
default/s/t) plus the matching dispatch in unified_attention.cpp and the
fmha_batch_prefill codegen hook.

Co-authored-by: Cursor <cursoragent@cursor.com>

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
 
 

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

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

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

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

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

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

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 <KernelVariant V,
 std::pair<bool, float> dispatch_page_size(const unified_attention_args& args,
                                           const stream_config& config)
 {
+    if(ua_force_runtime_ps_enabled())
+        return dispatch_one<V, DType, IsMask, 0>(args, config);
+
     if constexpr(V == KernelVariant::prefill_d128 ||
                  V == KernelVariant::prefill_d64)
     {
@@ -232,11 +250,31 @@ std::pair<bool, float> dispatch_page_size(const unified_attention_args& args,
             default: return dispatch_one<V, DType, IsMask, 0>(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<V, DType, IsMask, 128>(args, config);
+            default:  return dispatch_one<V, DType, IsMask, 0>(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<V, DType, IsMask, 0>(args, config);
     }
 }

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)