From 53b581e122e308d5d7ba5d02b1413511abb192b4 Mon Sep 17 00:00:00 2001 From: Po Yen Chen Date: Sun, 22 Sep 2024 20:05:58 +0800 Subject: [PATCH] Early return if seqlen_k=0 on group mode (#1524) [ROCm/composable_kernel commit: 770d2b77253b5bfbcc794d4133e7ecada63cdd44] --- .../ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp | 4 ++-- .../ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp index a392f0124d..281ddc07b7 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp @@ -215,8 +215,8 @@ struct BlockFmhaPipelineQRKSVS const auto num_total_loop = integer_divide_ceil(seqlen_k_end - seqlen_k_start, kN0); - // check early exit if masked and no work to do. - if constexpr(FmhaMask::IsMasking) + // check early exit if no work to do + if constexpr(FmhaMask::IsMasking || kPadSeqLenK) { if(num_total_loop <= 0) { diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp index 601aa62c5c..19f569c45b 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async.hpp @@ -268,7 +268,7 @@ struct BlockFmhaPipelineQRKSVSAsync const auto num_total_loop = integer_divide_ceil(seqlen_k_end - seqlen_k_start, kN0); - // check early exit + // check early exit if no work to do if constexpr(FmhaMask::IsMasking || kPadSeqLenK) { if(num_total_loop <= 0)