mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-29 19:28:33 +00:00
d9df801be2
Extends the fp8 paged ps128 work to bf16/fp16 prefill_d128 and removes the paged addressing overhead that made paged ~50% slower than contiguous. - Constexpr ps128 instances for bf16/fp16 (d128 prefill, mask + nmask). The dispatch in unified_attention.cpp now routes page_blk_size==128 to a constexpr instance for every dtype, so the per-tile div/mod/mul-by-128 strength-reduces to shifts/masks instead of falling into the PageSize=0 runtime-divide catch-all (which left real 32-bit integer divides in the bf16/fp16 address chain). - Decouple the single-page SRD rebase (kRebaseKSrd/kRebaseVSrd) from the scalar-promote gate. The rebase only needs the tile to fit in one page (kPageSize % kPageBlockSize == 0); it was wrongly gated behind NRepeat>=2, which excluded the trivial single-issue (NRepeat==1) tile -- bf16/fp16 d128 @ ps128 -- forcing it onto the per-lane multi-page fallback. Now single-issue tiles fold the page base into the SRD and hoist the per-lane scatter offsets; the paged 'addr' phase disappears from the ATT trace. - bf16/fp16 prefill_d128 quarters the N tile to 32 (kBf16QuarterBlockN) to drain the 2-byte score/P live set under the 256-VGPR ceiling (fp8's 1-byte tile fits at the halved N=64). - Route the multi-page K fallback (ps16/ps32) through the LDS-resident block-table cache (kKFallbackLds, on by default): LDS-resolving the per-lane block-table reads off the critical path recovers ~+10.5% at ps16/ps32. kKNContigLoad / kKMultiPageDedup added as default-off experimental levers. bf16 paged causal canonical (b1 sq=sk=75600 hq=hk=5 d128 bs128): 9.97 ms, 1.19x over Triton (was 10.5 ms / 1.13x). fp8 paged unchanged (1.64x). Full correctness matrix 263/263 PASS incl. all split-KV/causal/GQA regression fixtures. Co-authored-by: Cursor <cursoragent@cursor.com>