ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-13 18:51:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

epilogue switched to cshuffle

Browse Source
This commit is contained in:
aska-0096
2025-07-23 03:40:15 +00:00
parent 14e0ab70c6
commit 769fbb62d5
7 changed files with 49 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
example/ck_tile/01_fmha/CMakeLists.txt
View File

@@ -30,7 +30,7 @@ string(REPLACE ";" "," FMHA_FWD_APIS "${FMHA_FWD_ENABLE_APIS}")
set(FMHA_FWD_CODE_GEN_COMMON_ARGS
${CMAKE_CURRENT_LIST_DIR}/generate.py
--api ${FMHA_FWD_APIS}
# --filter @fmha_fwd_decode_d64_bf16_batch_b16x32x64x64x32x64_r1x1x1_r1x1x1_w16x16x32_w16x16x32_decode_qr_vr_nlogits_nbias_nmask_nlse_nsquant_npagedkv
# --filter @fmha_fwd_decode_d128_bf16_batch_b16x32x128x128x32x128_r1x1x1_r1x1x1_w16x16x32_w16x16x32_decode_qr_vr_npad_nlogits_nbias_nmask_nlse_nsquant_npagedkv
)
set(FMHA_BWD_CODE_GEN_COMMON_ARGS
${CMAKE_CURRENT_LIST_DIR}/generate.py

2
example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
View File

@@ -31,7 +31,7 @@ K0_MAX_SUBMAX_MAP = {
}
FMHA_FWD_KERNEL_HEADER = """// SPDX-License-Identifier: MIT
// Copyright (c) 2018-2024, Advanced Micro Devices, Inc. All rights reserved.\n
// Copyright (c) 2018-2025, Advanced Micro Devices, Inc. All rights reserved.\n
// auto generated by generate.py
#include "ck_tile/ops/fmha/block/variants.hpp"
#include "fmha_fwd.hpp"

25
example/ck_tile/01_fmha/codegen/ops/fmha_fwd_decode.py
View File

@@ -101,9 +101,30 @@ using fmha_pipeline = {F_pipeline}<
/// FIXME: use {F_spad}/{F_dvpad} as kPadM/kPadN parameters after solving
/// store_tile_raw() data corruption issue
using fmha_epilogue =
ck_tile::Default2DEpilogue<ck_tile::Default2DEpilogueProblem<typename FmhaFwdTypeConfig<{F_dtype}>::OaccDataType,
ck_tile::CShuffleEpilogue<ck_tile::CShuffleEpilogueProblem<
typename FmhaFwdTypeConfig<{F_dtype}>::PDataType,
typename FmhaFwdTypeConfig<{F_dtype}>::VDataType,
ck_tile::tuple<>,
typename FmhaFwdTypeConfig<{F_dtype}>::OaccDataType,
typename FmhaFwdTypeConfig<{F_dtype}>::ODataType,
false, false>>;
ck_tile::tuple<>,
ck_tile::tensor_layout::gemm::RowMajor,
ck_tile::element_wise::PassThrough,
{F_rm0}*{F_rn0}*64,
{F_bm0},
{F_bn1},
{F_rm1},
{F_rn1},
{F_wm1},
{F_wn1},
{F_wk1},
true,
ck_tile::integral_constant<ck_tile::memory_operation_enum,
ck_tile::memory_operation_enum::set>{{}}.value,
1,
true,
16/sizeof(typename FmhaFwdTypeConfig<{F_dtype}>::ODataType)>>;
static_assert(16/sizeof(typename FmhaFwdTypeConfig<{F_dtype}>::ODataType)==8);
using fmha_kernel =
ck_tile::FmhaFwdDecodeKernel<fmha_pipeline, fmha_epilogue>;

1
example/ck_tile/01_fmha/fmha_fwd.hpp
View File

@@ -6,6 +6,7 @@
#include "ck_tile/core.hpp"
#include "ck_tile/host/kernel_launch.hpp"
#include "ck_tile/ops/epilogue.hpp"
#include "ck_tile/ops/elementwise.hpp"
#include "ck_tile/ops/fmha.hpp"
#include "bias.hpp"