mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-29 19:28:33 +00:00
a3714e82cf
Eight files outside the UA scope had drifted onto this branch over time
via earlier commits whose subject lines explicitly carried no "CK-UA:"
prefix — they are independent fmha bug fixes and codegen additions that
do not touch any code path the unified_attention example, kernel or
pipeline actually compiles or includes.
Reverted to the merge-base content of:
include/ck_tile/ops/fmha/block/block_masking.hpp (-39 lines)
added by 6729989b9 "Fix FMHA split-KV for paged-KV with
page_block_size < kN0"
include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp (-2 lines)
added by ec2db01e4 "Fix fmha_fwd early-exit bug: seqlen_q <=
min_seqlen_q should be <"
example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
example/ck_tile/01_fmha/codegen/ops/fmha_pagedkv_prefill.py
example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py
example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
example/ck_tile/01_fmha/fmha_fwd_runner.hpp
example/ck_tile/01_fmha/mask.hpp
added by 63821af1f / cb6fb2802 / c5600bc8a / e5272603c / 10564b0c4
/ cd7ba6e2e / 07ba03bcb — split-KV decode tiles, codegen tweaks,
and a sliding-window mask fix, all in the 01_fmha example program
(a separate build target; the UA example lives in
42_unified_attention and pulls in zero 01_fmha sources).
These commits are still reachable from the branch's reflog and from
their original commit hashes; they should each be cherry-picked onto
their own branches and sent upstream as standalone fmha bug-fix PRs —
they look like clean fixes that upstream would welcome, but they don't
belong in the UA PR's scope.
Verified empirically: clean JIT rebuild of module_unified_attention
followed by both regression shapes pass at full perf
b=128/sk=16384/d=128/bf16 : 1.5152 ms, 5672 GB/s, PASS
b=1/sk=1M/d=128/bf16 nb=70k : 0.7677 ms, 5594 GB/s, PASS
matching the pre-revert numbers to within run-to-run noise.
Branch's shared-CK touch surface after this revert: tile_scatter_gather.hpp
(+152 from our async_load_raw_long method), load_tile.hpp (+21 from the
sister dispatcher), warp_gemm[.|_dispatcher.]hpp (+13 for the FP8 e4m3
small-tile registration), and the new amd_global_load_lds_raw.hpp file.
Down from 14 shared files to 4.
Co-authored-by: Cursor <cursoragent@cursor.com>