From c0e985d0756128ffcbf31f5d669ed72e9a68a38c Mon Sep 17 00:00:00 2001 From: juuso-oskari Date: Fri, 15 May 2026 10:18:39 +0000 Subject: [PATCH] CK-UA: document why per-issue SRD-rebase path was tried and dropped Replace the speculative TODO-style comment next to the K_mem_load / V_mem_load dispatch with a record of the actual experiment: we implemented async_load_tile_raw_rebased (buffer_load_dword_lds with a per-issue SRD whose 48-bit base absorbs the wave-uniform page offset), verified correctness on multiple big-cache decode shapes, and measured it against the existing async_load_tile_raw_long path on an isolated GPU. Rebased was at best tied with long and at worst ~6% slower (b=1 sk=1M d=64 GQA8: 2.46 ms vs 2.32 ms; b=8 sk=200k d=128 GQA8: 2.12 ms vs 2.02 ms). The workloads are compute / softmax bound, not K/V load bandwidth bound, so the buffer_load throughput edge never materialises, while the per-issue SRD construction adds real SGPR pressure. No functional change in this commit -- only the explanatory comment is updated so the next person who eyes the same idea finds the receipts before re-implementing. Co-authored-by: Cursor --- .../pipeline/unified_attention_pipeline.hpp | 21 ++++++++++++------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp index ee59d8c8bd..44b6e17a7a 100644 --- a/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp +++ b/include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp @@ -602,14 +602,19 @@ struct UnifiedAttentionPipeline // at the cost of lower throughput. // The branch is on a wave-uniform value, so no execution divergence. // - // For diagnostic purposes: the wave's N-position span within a - // single buffer_load instruction is `(LaneGroups-1)*NumWarps + 1`. - // When that's > the minimum page_size (≈16) the K-tile distribution - // touches multiple pages per issue, so the small-cache buffer_load - // path still works (per-lane voffsets fit while the cache ≤ 4 GB) - // but the per-issue SRD-rebase optimization (not implemented today) - // wouldn't be applicable — only `global_load_lds` works once the - // cache exceeds 4 GB. + // We tried a third "per-issue SRD rebase" path + // (`async_load_tile_raw_rebased`: buffer_load_dword_lds with a + // per-issue SRD whose 48-bit base absorbs the wave-uniform page + // offset, valid when WaveSpanInN ≤ runtime page_size). It was + // correct on every big-cache decode shape tested but came out at + // best tied with the long path and at worst ~6% slower (e.g. + // b=1 sk=1M d=64: 2.46 ms vs 2.32 ms; b=8 sk=200k d=128: 2.12 ms + // vs 2.02 ms — see git log for the full table). These workloads + // are compute / softmax bound, not K/V-load bandwidth bound, so + // the buffer_load vs global_load_lds throughput edge never + // materialises, while per-issue SRD construction adds real SGPR + // pressure. The rebased helper has been removed to keep the + // dispatch (and emitted kernel size) minimal. constexpr index_t KWaveSpanInN = (KDstrType::DstrEncode::hs_lengthss_[number<0>{}][number<1>{}] - 1) * KDstrType::DstrEncode::hs_lengthss_[number<0>{}][number<2>{}] +