Avoid seqlen_k=0 for kvcache

cgmillette · Aug 19, 2024 · 21c4df8 · 21c4df8
1 parent 3fb77a0
commit 21c4df8
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/example/ck_tile/01_fmha/utils.hpp b/example/ck_tile/01_fmha/utils.hpp
@@ -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