ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-12 18:17:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

CK-UA: document why per-issue SRD-rebase path was tried and dropped

Browse Source 
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 <cursoragent@cursor.com>
This commit is contained in:
juuso-oskari
2026-05-15 10:18:39 +00:00
parent 1f69421434
commit c0e985d075
1 changed files with 13 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
include/ck_tile/ops/unified_attention/pipeline/unified_attention_pipeline.hpp
View File

@@ -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>{}] +