diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index c92e677e6c..ba826c8fb3 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py @@ -114,8 +114,11 @@ template<> void fmha_fwd_splitkv_oneshot_(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a) {{ if constexpr({F_mode} == false) {{ // batch mode + // we don't check every seqlen_k values for kvcache + if (a.seqlen_k_ptr != nullptr) {{ + kernel_runner::run(s, a); // make sure F_bn0 is divisible by F_bk1 - if (a.seqlen_k % (a.num_splits * {F_bn0}) == 0) {{ + }} else if (a.seqlen_k % (a.num_splits * {F_bn0}) == 0) {{ kernel_runner::run(s, a); }} else {{ kernel_runner::run(s, a);