From 3b341e4a1641d28b3862189ca64a24945caebbd4 Mon Sep 17 00:00:00 2001 From: Qianfeng Date: Mon, 24 Nov 2025 21:31:48 +0800 Subject: [PATCH] Fix a bug for qr_ks_vs_async_trload pipeline (#3271) [ROCm/composable_kernel commit: 81042ea5747d3e1e4a71c3f327556f3fb0655d99] --- include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp | 2 +- .../fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp index cd1bfb031b..93b415b4ce 100644 --- a/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp +++ b/include/ck_tile/ops/fmha/kernel/fmha_fwd_kernel.hpp @@ -1641,7 +1641,7 @@ struct FmhaFwdKernel // 2. use more LDS, as we want better memory latency hiding // If SplitKV off, we don't expect Q data reused by different ThreadGroups, bypass the // cache - constexpr bool PrefillCase = FmhaPipeline::kM0 >= 128; + constexpr bool PrefillCase = FmhaPipeline::kM0 > 64; // divide problem const auto [i_tile_m, i_tile_n, i_nhead, i_batch] = GetTileIndex(kargs); diff --git a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp index 08fc42a471..dce9583fc1 100644 --- a/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp +++ b/include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs_async_trload.hpp @@ -37,7 +37,7 @@ struct BlockFmhaPipelineQRKSVSAsyncTrload using VLayout = remove_cvref_t; static constexpr bool kQLoadOnce = true; // if q_tile load whole block length (hdim) at once static_assert(kQLoadOnce == Policy::QLoadOnce); - static constexpr bool kKLoadOnce = BlockFmhaShape::kM0 >= 64; + static constexpr bool kKLoadOnce = BlockFmhaShape::kM0 > 64; static constexpr index_t kBlockSize = Problem::kBlockSize;