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

Add QK swizzle option

Browse Source
This commit is contained in:
MHYang
2025-04-29 16:58:15 +00:00
parent c204cdc382
commit 1b0885c207
8 changed files with 35 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
example/ck_tile/99_toy_example/03_flash_attention_fwd/CMakeLists.txt
View File

@@ -7,6 +7,8 @@ add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL flash_attention_fwd.cpp)
target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
# list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -v --save-temps -Wno-gnu-line-marker)
# NOTE: we turn off undefined-func-template to let source compile without explicit declare function specializations
list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS -Wno-undefined-func-template -Wno-float-equal)
@@ -16,6 +18,12 @@ if(ENABLE_TOY_FA_FWD_OPT)
target_compile_definitions(${EXAMPLE_REDUCE} PRIVATE TOY_FA_FWD_OPT)
endif()
option(ENABLE_TOY_FA_FWD_QK_SWIZZLE "Enable toy FA fwd QK swizzle" OFF)
if(ENABLE_TOY_FA_FWD_QK_SWIZZLE)
message("Compiling with toy FA fwd QK swizzle")
target_compile_definitions(${EXAMPLE_REDUCE} PRIVATE TOY_FA_FWD_QK_SWIZZLE)
endif()
target_compile_options(${EXAMPLE_REDUCE} PRIVATE ${EXAMPLE_REDUCE_COMPILE_OPTIONS})
# TODO: we have to turn off this global prop, otherwise the progress bar generated

5
example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
View File

@@ -23,7 +23,12 @@ struct BlockGemmARegBSmemCRegV1DefaultPolicy
}
else if constexpr(kM0 == 128)
{
#if !defined(TOY_FA_FWD_QK_SWIZZLE)
return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1);
#else
#pragma message("Enable toy FA fwd QK swizzle")
return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
#endif
}
else
{

4
example/ck_tile/99_toy_example/03_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
View File

@@ -23,7 +23,11 @@ struct BlockGemmARegBSmemCRegV1K8Policy
}
else if constexpr(kM0 == 128)
{
#if !defined(TOY_FA_FWD_QK_SWIZZLE)
return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
#else
return make_tuple(WarpGemmMfmaF16F16F32M32N32K16SwizzleBTransposedCDistribution{}, 4, 1);
#endif
}
else
{

4
example/ck_tile/99_toy_example/03_flash_attention_fwd/flash_attention_fwd_impl.hpp
View File

@@ -63,7 +63,11 @@ struct FlashAttentionFwdImpl
{
constexpr index_t kNPerBlock = kN1PerBlock;
constexpr index_t kKPerBlock = kK1PerBlock;
#if !defined(TOY_FA_FWD_QK_SWIZZLE)
constexpr index_t kKPack = 4;
#else
constexpr index_t kKPack = 8;
#endif
constexpr auto dataTypeSize = sizeof(VDataType);
constexpr auto NLdsLayer =

1
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
View File

@@ -6,6 +6,7 @@ if(FLASH_ATTENTION_FWD_ENABLE_APIS STREQUAL "all")
endif()
option(TOY_FA_FWD_OPT "Enable toy flash attention forward optimization" ON)
option(TOY_FA_FWD_QK_SWIZZLE "Enable toy flash attention forward QK swizzle" OFF)
execute_process(
COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py

5
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_default_policy.hpp
View File

@@ -23,7 +23,12 @@ struct BlockGemmARegBSmemCRegV1DefaultPolicy
}
else if constexpr(kM0 == 128)
{
#if !defined(TOY_FA_FWD_QK_SWIZZLE)
return make_tuple(WarpGemmMfmaF16F16F32M32N32K8TransposedCDistribution{}, 4, 1);
#else
#pragma message("Enable toy FA fwd QK swizzle")
return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
#endif
}
else
{

4
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/block_gemm_areg_bsmem_creg_v1_iteratek_policy.hpp
View File

@@ -23,7 +23,11 @@ struct BlockGemmARegBSmemCRegV1K8Policy
}
else if constexpr(kM0 == 128)
{
#if !defined(TOY_FA_FWD_QK_SWIZZLE)
return make_tuple(WarpGemmMfmaF16F16F32M32N32K16TransposedCDistribution{}, 4, 1);
#else
return make_tuple(WarpGemmMfmaF16F16F32M32N32K16SwizzleBTransposedCDistribution{}, 4, 1);
#endif
}
else
{

4
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd_impl.hpp
View File

@@ -63,7 +63,11 @@ struct FlashAttentionFwdImpl
{
constexpr index_t kNPerBlock = kN1PerBlock;
constexpr index_t kKPerBlock = kK1PerBlock;
#if !defined(TOY_FA_FWD_QK_SWIZZLE)
constexpr index_t kKPack = 4;
#else
constexpr index_t kKPack = 8;
#endif
constexpr auto dataTypeSize = sizeof(VDataType);
constexpr auto NLdsLayer =