ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-06-07 08:15:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid seqlen_k=0 for kvcache

Browse Source
This commit is contained in:
PoYen, Chen
2024-08-19 07:13:57 +00:00
parent 3fb77a0ebb
commit 21c4df89e4
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
example/ck_tile/01_fmha/utils.hpp
View File

@@ -149,7 +149,10 @@ decode_seqlen(mode_enum mode,
if(use_kvcache)
{
std::vector<ck_tile::index_t> seqlen_ks(batch);
randints(seqlen_ks.begin(), seqlen_ks.end(), seqlen_k_min, seqlen_k_max, seed);
// generate random numbers in range [1, seqlen_k_max] if given seqlen_k_min = 0
ck_tile::index_t real_seqlen_k_min =
std::min(std::max(1, seqlen_k_min), seqlen_k_max);
randints(seqlen_ks.begin(), seqlen_ks.end(), real_seqlen_k_min, seqlen_k_max, seed);
return seqlen_ks;
}
else