ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-06 05:55:39 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix wrong knew/vew appending logic on host

Browse Source
This commit is contained in:
PoYen, Chen
2024-08-20 23:58:47 +00:00
parent 3850028bca
commit 73378ff95d
1 changed files with 4 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
example/ck_tile/01_fmha/fmha_fwd.cpp
View File

@@ -1198,12 +1198,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
real_knew_host_ref = &knew_host_ref_ro.value();
}
const std::size_t knew_start = real_seqlen_k - seqlen_knew;
k_host_ref.ForEach([&](auto& self, auto i) {
if(knew_start <= i[1])
{
self(i) = (*real_knew_host_ref)(i[0], i[1] - knew_start, i[2]);
}
(*real_knew_host_ref).ForEach([&](auto& self, auto i) {
k_host_ref(i[0], i[1] + cache_seqlen_ks[wb], i[2]) = self(i);
});
}
#endif
@@ -1264,12 +1260,8 @@ bool run(const ck_tile::ArgParser& arg_parser)
else vnew_host_ref.ForEach([&](auto& self, auto i) { self(i) = vnew_host(wb, i[1], i[0] / nr, i[2]); });
}
const std::size_t knew_start = real_seqlen_k - seqlen_knew;
v_host_ref.ForEach([&](auto& self, auto i) {
if(knew_start <= i[2])
{
self(i) = vnew_host_ref(i[0], i[1], i[2] - knew_start);
}
vnew_host_ref.ForEach([&](auto& self, auto i) {
v_host_ref(i[0], i[1], i[2] + cache_seqlen_ks[wb]) = self(i);
});
}
#endif