composable_kernel icon indicating copy to clipboard operation
composable_kernel copied to clipboard

[CK_TILE] Add appendkv kernel to support mha with kvcache

Open poyenc opened this issue 1 year ago • 1 comments

Add new fmha_fwd_appendkv() API which runs ahead the fmha_fwd()/fmha_fwd_splitkv() API. The fmha_fwd_appendkv() + fmha_fwd()/fmha_fwd_splitkv() combination implement the functionality of mha_fwd_kvcache() in FA 2.5 (without paged-kvcache part)

poyenc avatar Jul 14 '24 22:07 poyenc

the paged-kvcache will be added to this PR.

poyenc avatar Jul 30 '24 20:07 poyenc

All the test has pass (MI200 + MI300 @ ROCm6.1) in flash attention https://github.com/ROCm/flash-attention/pull/74

rocking5566 avatar Aug 27 '24 15:08 rocking5566