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 9e9cce5400..0836fbfce3 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 @@ -578,6 +578,9 @@ struct BlockFmhaPipelineQRKSVS if constexpr(kHasDropout) { + // K and dropout use the same address in LDS, finish loading from k_lds_window by + // gemm_0 to reuse LDS. + block_sync_lds(); dropout.template Run( smem_ptr, seqlen_k_start + i_total_loops * kN0, p_compute, randval_dram_window); }