fix epilogue problem

0d93f4a0 · danyao12 · edb77c8d · 0d93f4a0
Commit 0d93f4a0 authored Jul 22, 2024 by danyao12
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 4 deletions

example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py +4 -4

No files found.
--- a/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py
+++ b/example/ck_tile/01_fmha/codegen/ops/fmha_bwd.py
@@ -94,14 +94,14 @@ using fmha_bwd_pipeline_{F_idx} = {F_pipeline}<fmha_bwd_pipeline_problem_{F_idx}
 using fmha_bwd_dk_epilogue_{F_idx} = ck_tile::Default2DEpilogue<
    ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
                                      typename FmhaBwdTypeConfig<{F_dtype}>::KGradDataType,
-                                      false,
-                                      false>>;
+                                      {F_skpad},
+                                      {F_dpad}>>;

 using fmha_bwd_dv_epilogue_{F_idx} = ck_tile::Default2DEpilogue<
    ck_tile::Default2DEpilogueProblem<typename FmhaBwdTypeConfig<{F_dtype}>::AccDataType,
                                      typename FmhaBwdTypeConfig<{F_dtype}>::VGradDataType,
-                                      false,
-                                      false>>;
+                                      {F_skpad},
+                                      {F_dvpad}>>;

 using fmha_bwd_dq_dk_dv_kernel_{F_idx} =
    ck_tile::FmhaBwdDQDKDVKernel<fmha_bwd_pipeline_{F_idx},