ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-05-02 12:41:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix synchronization issue in fwd qr pipeline with dropout (#3135)

Browse Source 
BlockFmhaPipelineQRKSVS reuses LDS for K and dropout so there must be
block_sync_lds between loading k_lds_window by gemm_0 and storing
dropout randval.
This commit is contained in:
Anton Gorenko
2025-10-31 21:44:52 +05:00
committed by GitHub
parent 5ed2046bee
commit e9596228ff
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
include/ck_tile/ops/fmha/pipeline/block_fmha_pipeline_qr_ks_vs.hpp
View File

@@ -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<decltype(gemm_0), SMPLComputeDataType, RandValOutputDataType>(
smem_ptr, seqlen_k_start + i_total_loops * kN0, p_compute, randval_dram_window);
}