From 25db1339265fa020d457e13d8440786d647fcc23 Mon Sep 17 00:00:00 2001 From: danyao12 Date: Mon, 5 Aug 2024 17:18:10 +0000 Subject: [PATCH] dpad same as dvpad for flash attention integration --- example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py | 1 + 1 file changed, 1 insertion(+) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py index d1929ab89f..096394c0c9 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py @@ -496,6 +496,7 @@ def get_bwd_dq_dk_dv_blobs(kernel_filter : Optional[str], receipt, mask_impl) -> cond = dtype in ['fp16', 'bf16'] cond &= bias in ['no', 'alibi'] cond &= dropout in ['no', 'dropout_wg32', 'dropout_wg16'] + cond &= dpad == dvpad if not cond: continue if receipt == 3: