diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt index 93a1364b37..cc9873eeb3 100755 --- a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt @@ -1,19 +1,59 @@ +set(FLASH_ATTENTION_FWD_KNOWN_APIS "fwd") +set(FLASH_ATTENTION_FWD_ENABLE_APIS "fwd" CACHE STRING + "semicolon-separated list of APIs to generate (${FLASH_ATTENTION_FWD_KNOWN_APIS}) & link, or \"all\".") +if(FLASH_ATTENTION_FWD_ENABLE_APIS STREQUAL "all") + set(FLASH_ATTENTION_FWD_ENABLE_APIS ${FLASH_ATTENTION_FWD_KNOWN_APIS}) +endif() + +execute_process( + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py + --api ${FLASH_ATTENTION_FWD_ENABLE_APIS} + --working_path ${CMAKE_CURRENT_BINARY_DIR} + --list_blobs + RESULT_VARIABLE ret +) +if(ret AND NOT ret EQUAL 0) + message(FATAL_ERROR "Failed to list Flash Attention kernels via Python. ${ret}") +endif() + +file(STRINGS ${CMAKE_CURRENT_BINARY_DIR}/flash_attention_fwd_blobs.txt FLASH_ATTENTION_FWD_GEN_BLOBS) + +add_custom_command( + OUTPUT ${FLASH_ATTENTION_FWD_GEN_BLOBS} + COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_LIST_DIR}/generate.py + --api ${FLASH_ATTENTION_FWD_ENABLE_APIS} + --working_path ${CMAKE_CURRENT_BINARY_DIR} + --gen_blobs +) + set(EXAMPLE_REDUCE "codegen_basic_flash_attention_fwd") -# not using add_example_executable() to add this target, since we don't want this to have -# to be included in "make all/install/check" message("adding example ${EXAMPLE_REDUCE}") -add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL flash_attention_fwd.cpp) -target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR}) +add_executable(${EXAMPLE_REDUCE} + EXCLUDE_FROM_ALL + flash_attention_fwd.cpp +) + +target_include_directories(${EXAMPLE_REDUCE} + PRIVATE + ${CMAKE_CURRENT_LIST_DIR} +) + +target_sources(${EXAMPLE_REDUCE} PRIVATE ${FLASH_ATTENTION_FWD_GEN_BLOBS}) + +message("FLASH_ATTENTION_FWD_GEN_BLOBS = ${FLASH_ATTENTION_FWD_GEN_BLOBS}") + + set(EXAMPLE_REDUCE_COMPILE_OPTIONS) +list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS + -Wno-undefined-func-template + -Wno-float-equal + --offload-compress +) -# 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) +target_compile_options(${EXAMPLE_REDUCE} + PRIVATE + ${EXAMPLE_REDUCE_COMPILE_OPTIONS} +) -target_compile_options(${EXAMPLE_REDUCE} PRIVATE ${EXAMPLE_REDUCE_COMPILE_OPTIONS}) - -# TODO: we have to turn off this global prop, otherwise the progress bar generated -# by cmake will print too many files, execvp: /bin/sh: Argument list too long -# however, this property may affect global -# TODO: consider codegen a makefile by us -set_property(GLOBAL PROPERTY RULE_MESSAGES OFF) +set_property(GLOBAL PROPERTY RULE_MESSAGES OFF) \ No newline at end of file diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp index 77bfbd83cc..e3f4fde0bb 100644 --- a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/flash_attention_fwd.cpp @@ -131,7 +131,7 @@ int main(int argc, char* argv[]) }; float ave_time = ck_tile::flash_attention_fwd -struct flash_attention_fwd_traits_ -{ - using SaccDataType = ck_tile::remove_cvref_t; - using SMPLComputeDataType = ck_tile::remove_cvref_t; - using PDataType = ck_tile::remove_cvref_t; - using OaccDataType = ck_tile::remove_cvref_t; +// // TODO: fwd_api.cpp +// template +// struct flash_attention_fwd_traits_ +// { +// using SaccDataType = ck_tile::remove_cvref_t; +// using SMPLComputeDataType = ck_tile::remove_cvref_t; +// using PDataType = ck_tile::remove_cvref_t; +// using OaccDataType = ck_tile::remove_cvref_t; - static constexpr index_t kBlockSize = kBlockSize_; - static constexpr index_t kHeadDim = kHeadDim_; - static constexpr index_t kM0PerBlock = kM0PerBlock_; - static constexpr index_t kN0PerBlock = kN0PerBlock_; - static constexpr index_t kK0PerBlock = kK0PerBlock_; - static constexpr index_t kN1PerBlock = kN1PerBlock_; - static constexpr index_t kK1PerBlock = kK1PerBlock_; +// static constexpr index_t kBlockSize = kBlockSize_; +// static constexpr index_t kHeadDim = kHeadDim_; +// static constexpr index_t kM0PerBlock = kM0PerBlock_; +// static constexpr index_t kN0PerBlock = kN0PerBlock_; +// static constexpr index_t kK0PerBlock = kK0PerBlock_; +// static constexpr index_t kN1PerBlock = kN1PerBlock_; +// static constexpr index_t kK1PerBlock = kK1PerBlock_; - static constexpr ck_tile::index_t kWarpPerCu = 8; // 2 warps per SIMD - static constexpr ck_tile::index_t kWarpPerBlock = kBlockSize / warpSize; - static constexpr ck_tile::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock; -}; +// static constexpr ck_tile::index_t kWarpPerCu = 8; // 2 warps per SIMD +// static constexpr ck_tile::index_t kWarpPerBlock = kBlockSize / warpSize; +// static constexpr ck_tile::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock; +// }; -// TODO: fwd_api.cpp, fwd_common.cpp -template -using traits_ = flash_attention_fwd_traits_; +// // TODO: fwd_api.cpp, fwd_common.cpp +// template +// using traits_ = flash_attention_fwd_traits_; +// // fw_api.cpp +// // Note: this internal API only declare, not define here, otherwise will block `make -j` +// template +// float flash_attention_fwd_(const FlashAttnArgs& a, +// const ck_tile::stream_config& stream_config); -// Note: this internal API only declare, not define here, otherwise will block `make -j` -template -float flash_attention_fwd_(const FlashAttnArgs& a, - const ck_tile::stream_config& stream_config); - -// TODO: fwd_common.cpp -template -float flash_attention_fwd_(const FlashAttnArgs& a, - const ck_tile::stream_config& stream_config) { - using SaccDataType = typename Traits_::SaccDataType; - using SMPLComputeDataType = typename Traits_::SMPLComputeDataType; - using PDataType = typename Traits_::PDataType; - using OaccDataType = typename Traits_::OaccDataType; +// // TODO: fwd_common.cpp +// template +// float flash_attention_fwd_(const FlashAttnArgs& a, +// const ck_tile::stream_config& stream_config) { +// using SaccDataType = typename Traits_::SaccDataType; +// using SMPLComputeDataType = typename Traits_::SMPLComputeDataType; +// using PDataType = typename Traits_::PDataType; +// using OaccDataType = typename Traits_::OaccDataType; - index_t kGridSize = a.Batch * (a.M0 / Traits_::kM0PerBlock) * (a.N1 / Traits_::kN1PerBlock); +// index_t kGridSize = a.Batch * (a.M0 / Traits_::kM0PerBlock) * (a.N1 / Traits_::kN1PerBlock); - std::cout << "grid size " << kGridSize << std::endl; +// std::cout << "grid size " << kGridSize << std::endl; + +// return ck_tile::launch_kernel(stream_config, +// ck_tile::make_kernel( +// ck_tile::FlashAttentionFwd{}, +// kGridSize, +// Traits_::kBlockSize, +// 0, +// a.q_ptr, +// a.k_ptr, +// a.v_ptr, +// a.o_ptr, +// a.M0, +// a.N0, +// a.K0, +// a.N1, +// a.Batch, +// a.strideQ, // StrideQ +// a.strideK, // StrideK +// a.strideV, // StrideV +// a.strideO, // StrideO +// a.batchStrideQ, // BatchStrideQ +// a.batchStrideK, // BatchStrideK +// a.batchStrideV, // BatchStrideV +// a.batchStrideO)); // BatchStrideO +// } + +// // TODO: change to only declare +// // TODO: fwd_api.cpp +// template +// float flash_attention_fwd(const FlashAttnArgs& a, +// const ck_tile::stream_config& stream_config) { +// constexpr ck_tile::index_t kM0PerBlock = 128; +// constexpr ck_tile::index_t kN0PerBlock = 128; +// constexpr ck_tile::index_t kK0PerBlock = 32; +// constexpr ck_tile::index_t kN1PerBlock = 128; +// constexpr ck_tile::index_t kK1PerBlock = 32; + +// constexpr ck_tile::index_t kBlockSize = 256; +// constexpr ck_tile::index_t kHeadDim = 128; + +// return flash_attention_fwd_> +// (a, stream_config); + +// } - return ck_tile::launch_kernel(stream_config, - ck_tile::make_kernel( - ck_tile::FlashAttentionFwd{}, - kGridSize, - Traits_::kBlockSize, - 0, - a.q_ptr, - a.k_ptr, - a.v_ptr, - a.o_ptr, - a.M0, - a.N0, - a.K0, - a.N1, - a.Batch, - a.strideQ, // StrideQ - a.strideK, // StrideK - a.strideV, // StrideV - a.strideO, // StrideO - a.batchStrideQ, // BatchStrideQ - a.batchStrideK, // BatchStrideK - a.batchStrideV, // BatchStrideV - a.batchStrideO)); // BatchStrideO -} // TODO: change to only declare // TODO: fwd_api.cpp @@ -274,33 +316,7 @@ template float flash_attention_fwd(const FlashAttnArgs& a, - const ck_tile::stream_config& stream_config) { - constexpr ck_tile::index_t kM0PerBlock = 128; - constexpr ck_tile::index_t kN0PerBlock = 128; - constexpr ck_tile::index_t kK0PerBlock = 32; - constexpr ck_tile::index_t kN1PerBlock = 128; - constexpr ck_tile::index_t kK1PerBlock = 32; + const stream_config& stream_config); - constexpr ck_tile::index_t kBlockSize = 256; - constexpr ck_tile::index_t kHeadDim = 128; - - return flash_attention_fwd_> - (a, stream_config); - -} } // namespace ck_tile diff --git a/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/generate.py b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/generate.py new file mode 100644 index 0000000000..4a36179dbe --- /dev/null +++ b/example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/generate.py @@ -0,0 +1,628 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +import argparse +from enum import IntEnum +from pathlib import Path +import sys +from typing import List, Optional, Any +import functools +import itertools +import copy +from dataclasses import dataclass + +# def get_if_str(idx, total, last_else=True): +# if idx == 0: +# return 'if' +# elif idx < total - 1: +# return 'else if' +# else: +# return 'else' if last_else else 'else if' + +def get_if_str(size_, total, last_else=True): + if size_ == "small": + return 'if' + else: + return 'else if' + +DATA_TYPE_MAP = {'fp32': 'float', + 'fp16': 'ck_tile::half_t', + 'bf16': 'ck_tile::bf16_t'} + +def BOOL_MAP(b_) -> str: + return 'true' if b_ else 'false' + +class FlashAttentionFwdCodegen: + API_TRAITS_DEFINE = """ + +template +struct flash_attention_fwd_traits_ +{ + using SaccDataType = ck_tile::remove_cvref_t; + using SMPLComputeDataType = ck_tile::remove_cvref_t; + using PDataType = ck_tile::remove_cvref_t; + using OaccDataType = ck_tile::remove_cvref_t; + + static constexpr index_t kBlockSize = kBlockSize_; + static constexpr index_t kHeadDim = kHeadDim_; + static constexpr index_t kM0PerBlock = kM0PerBlock_; + static constexpr index_t kN0PerBlock = kN0PerBlock_; + static constexpr index_t kK0PerBlock = kK0PerBlock_; + static constexpr index_t kN1PerBlock = kN1PerBlock_; + static constexpr index_t kK1PerBlock = kK1PerBlock_; + + static constexpr ck_tile::index_t kWarpPerCu = 8; // 2 warps per SIMD + static constexpr ck_tile::index_t kWarpPerBlock = kBlockSize / warpSize; + static constexpr ck_tile::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock; +}; + +template +using traits_ = flash_attention_fwd_traits_; +""" + +# API_COMMON_HEADER = """ +# // SPDX-License-Identifier: MIT +# // Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved. + +# #include +# #include "flash_attention_fwd.hpp" +# #include + +# #pragma once + +# using S = ck_tile::stream_config; +# using A = FlashAttnArgs; + +# {F_traits_define} + +# template +# float flash_attention_fwd_(const FlashAttnArgs& a, +# const ck_tile::stream_config& stream_config) {{ +# using SaccDataType = typename Traits_::SaccDataType; +# using SMPLComputeDataType = typename Traits_::SMPLComputeDataType; +# using PDataType = typename Traits_::PDataType; +# using OaccDataType = typename Traits_::OaccDataType; + +# index_t kGridSize = a.Batch * (a.M0 / Traits_::kM0PerBlock) * (a.N1 / Traits_::kN1PerBlock); + +# if(stream_config.log_level_ > 0) +# std::cout << ", " << "FlashAttentionFwd<" << Traits_::kBlockSize << "," << Traits_::kHeadDim << ">" << std::flush; + +# return ck_tile::launch_kernel(stream_config, +# ck_tile::make_kernel( +# ck_tile::FlashAttentionFwd{{}}, +# kGridSize, +# Traits_::kBlockSize, +# 0, +# a.q_ptr, +# a.k_ptr, +# a.v_ptr, +# a.o_ptr, +# a.M0, +# a.N0, +# a.K0, +# a.N1, +# a.Batch, +# a.strideQ, // StrideQ +# a.strideK, // StrideK +# a.strideV, // StrideV +# a.strideO, // StrideO +# a.batchStrideQ, // BatchStrideQ +# a.batchStrideK, // BatchStrideK +# a.batchStrideV, // BatchStrideV +# a.batchStrideO)); // BatchStrideO +# }} +# """ + + API_BASE = """ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include +#include "flash_attention_fwd.hpp" + +namespace ck_tile {{ + +{F_traits_define} + +// Note: this internal API only declare, not define here, otherwise will block `make -j` +template +float flash_attention_fwd_(const FlashAttnArgs& a, + const ck_tile::stream_config& stream_config); + +template +float flash_attention_fwd(const FlashAttnArgs& a, + const ck_tile::stream_config& stream_config) {{ + float r = -1; +{F_dispatch} + return r; +}} + +template float flash_attention_fwd( + const FlashAttnArgs&, + const ck_tile::stream_config&); + +}} +""" + +# API_PER_DTYPE = """ {F_if}(std::is_same_v && std::is_same_v && std::is_same_v && std::is_same_v) {{ +# {F_per_size_case} +# }} +# """ +# API_PER_SIZE_CASE = """ {F_if} {F_SIZE_COND} {{ +# {F_inner_dispatch} +# }} +# """ + API_INNER_CASE = """ {F_if} {F_VEC_COND} + r = flash_attention_fwd_>(a, stream_config); +""" + + INSTANCE_BASE = """ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "flash_attention_fwd_api_common.hpp" + +namespace ck_tile { +// clang-format off +// +{F_instance_def} +// clang-format on + +} +""" + + def __init__(self, working_path, kernel_filter): + self.working_path = working_path + self.kernel_filter = kernel_filter + + @dataclass + class h_traits: + F_SaccDataType: str + F_SMPLComputeDataType: str + F_PDataType: str + F_OaccDataType: str + F_kBlockSize: int + F_kHeadDim: int + F_kM0PerBlock: int + F_kN0PerBlock: int + F_kK0PerBlock: int + F_kN1PerBlock: int + F_kK1PerBlock: int + + @property + def trait_name(self) -> str: + return (f"{DATA_TYPE_MAP[self.F_SaccDataType]}, " + f"{DATA_TYPE_MAP[self.F_SMPLComputeDataType]}, " + f"{DATA_TYPE_MAP[self.F_PDataType]}, " + f"{DATA_TYPE_MAP[self.F_OaccDataType]}, " + f"{self.F_kBlockSize}, {self.F_kHeadDim}, " + f"{self.F_kM0PerBlock}, {self.F_kN0PerBlock}, {self.F_kK0PerBlock}, " + f"{self.F_kN1PerBlock}, {self.F_kK1PerBlock}") + + @property + def def_name(self) -> str: + return (f"template float flash_attention_fwd_<{DATA_TYPE_MAP['fp16']}, " + f"{DATA_TYPE_MAP['fp16']}, {DATA_TYPE_MAP['fp16']}, {DATA_TYPE_MAP['fp16']}, " + f"traits_<{self.trait_name}>>(const FlashAttnArgs<{DATA_TYPE_MAP['fp16']}, " + f"{DATA_TYPE_MAP['fp16']}, {DATA_TYPE_MAP['fp16']}, {DATA_TYPE_MAP['fp16']}>&, " + "const ck_tile::stream_config&);") + + @dataclass + class h_instance: + F_DataTypePair: str # "q,k,v,o" + F_SizeCategory: str # "small", "medium", "large" + instance_list: List[Any] # List[h_traits] + + INSTANCE_BASE = """ +// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. + +#include "flash_attention_fwd_api_common.hpp" + +namespace ck_tile {{ +// clang-format off +// +{F_instance_def} +// clang-format on +}} +""" + + @property + def name(self) -> str: + q_type, k_type, v_type, o_type = self.F_DataTypePair.split(',') + dtype_str = f"{q_type}_{k_type}_{v_type}_{o_type}" + return f"flash_attention_fwd_{dtype_str}_{self.F_SizeCategory}" + + @property + def content(self) -> str: + instance_defs = '\n'.join(ins.def_name for ins in self.instance_list) + return self.INSTANCE_BASE.format(F_instance_def=instance_defs) + + @property + def name_api(self) -> str: + return "flash_attention_fwd_api" + + @property + def name_common_header(self) -> str: + return "flash_attention_fwd_api_common" + + @property + def content_common_header(self) -> str: + return f"""// SPDX-License-Identifier: MIT +// Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved. +#pragma once + +#include "flash_attention_fwd.hpp" + +namespace ck_tile {{ + +template +struct flash_attention_fwd_traits_ +{{ + using SaccDataType = ck_tile::remove_cvref_t; + using SMPLComputeDataType = ck_tile::remove_cvref_t; + using PDataType = ck_tile::remove_cvref_t; + using OaccDataType = ck_tile::remove_cvref_t; + + static constexpr index_t kBlockSize = kBlockSize_; + static constexpr index_t kHeadDim = kHeadDim_; + static constexpr index_t kM0PerBlock = kM0PerBlock_; + static constexpr index_t kN0PerBlock = kN0PerBlock_; + static constexpr index_t kK0PerBlock = kK0PerBlock_; + static constexpr index_t kN1PerBlock = kN1PerBlock_; + static constexpr index_t kK1PerBlock = kK1PerBlock_; + + static constexpr ck_tile::index_t kWarpPerCu = 8; // 2 warps per SIMD + static constexpr ck_tile::index_t kWarpPerBlock = kBlockSize / warpSize; + static constexpr ck_tile::index_t kBlockPerCu = kWarpPerCu / kWarpPerBlock; +}}; + + +template +using traits_ = flash_attention_fwd_traits_; + + +template +float flash_attention_fwd_(const FlashAttnArgs& a, + const ck_tile::stream_config& stream_config) {{ + using SaccDataType = typename Traits_::SaccDataType; + using SMPLComputeDataType = typename Traits_::SMPLComputeDataType; + using PDataType = typename Traits_::PDataType; + using OaccDataType = typename Traits_::OaccDataType; + + index_t kGridSize = a.Batch * (a.M0 / Traits_::kM0PerBlock) * (a.N1 / Traits_::kN1PerBlock); + + if(stream_config.log_level_ > 0) + std::cout << ", " << "FlashAttentionFwd<" << Traits_::kBlockSize << "," << Traits_::kHeadDim << ">" << std::flush; + + return ck_tile::launch_kernel(stream_config, + ck_tile::make_kernel( + ck_tile::FlashAttentionFwd{{}}, + kGridSize, + Traits_::kBlockSize, + 0, + a.q_ptr, + a.k_ptr, + a.v_ptr, + a.o_ptr, + a.M0, + a.N0, + a.K0, + a.N1, + a.Batch, + a.strideQ, // StrideQ + a.strideK, // StrideK + a.strideV, // StrideV + a.strideO, // StrideO + a.batchStrideQ, // BatchStrideQ + a.batchStrideK, // BatchStrideK + a.batchStrideV, // BatchStrideV + a.batchStrideO)); // BatchStrideO +}} +}} +""" + def content_api(self, args) -> str: + # Sort based on dtype + t_dtype_dict = {} + blobs = self.get_blobs(args) + + for blob in blobs: + if blob.F_DataTypePair not in t_dtype_dict: + t_dtype_dict[blob.F_DataTypePair] = {} + if blob.F_SizeCategory not in t_dtype_dict[blob.F_DataTypePair]: + t_dtype_dict[blob.F_DataTypePair][blob.F_SizeCategory] = [] + t_dtype_dict[blob.F_DataTypePair][blob.F_SizeCategory].append(blob) + + d_str = '' + for i_d, dtype_ in enumerate(t_dtype_dict): + blob_per_t = t_dtype_dict[dtype_] + size_str = '' + + for i_size, size_ in enumerate(blob_per_t): + blob_per_size = blob_per_t[size_] + inner_str = "" + + for i_b, b_ in enumerate(blob_per_size): + for i_ins, ins in enumerate(b_.instance_list): + idx_in_size = i_b * len(b_.instance_list) + i_ins + len_in_size = sum(len(b.instance_list) for b in blob_per_size) + + size_cond = "" + if size_ == "small": + size_cond = "(a.M0 < 2048 && a.N0 < 2048)" + elif size_ == "medium": + size_cond = "(a.M0 >= 2048 && a.N0 >= 2048 && a.M0 < 4096 && a.N0 < 4096)" + else: # large + size_cond = "(a.M0 >= 4096 || a.N0 >= 4096)" + + inner_str += self.API_INNER_CASE.format( + # F_if=get_if_str(idx_in_size, len_in_size, False), + F_if=get_if_str(size_, len_in_size, False), + F_VEC_COND=size_cond, + F_trait_name=ins.trait_name + ) + + # size_str += self.API_PER_SIZE_CASE.format( + # F_if=get_if_str(i_size, len(blob_per_t)), + # F_SIZE_COND=size_cond, + # F_inner_dispatch=inner_str + # ) + size_str += inner_str + + # q_type, k_type, v_type, o_type = dtype_.split(',') + # d_str += self.API_PER_DTYPE.format( + # F_if=get_if_str(i_d, len(t_dtype_dict)), + # F_q_type=DATA_TYPE_MAP[q_type], + # F_k_type=DATA_TYPE_MAP[k_type], + # F_v_type=DATA_TYPE_MAP[v_type], + # F_o_type=DATA_TYPE_MAP[o_type], + # F_per_size_case=size_str + # ) + d_str += size_str + + api_base = self.API_BASE.format( + F_traits_define=self.API_TRAITS_DEFINE, + F_dispatch=d_str + ) + return api_base + + def get_blobs(self, args): + h_traits = self.h_traits + h_instance = self.h_instance + + # Define kernel configurations for different size categories + trait_dict = { + "small": [ + h_traits('fp32', 'fp32', 'fp32', 'fp32', 128, 128, 128, 128, 32, 128, 32), + # h_traits('fp32', 'fp32', 'fp32', 'fp32', 256, 128, 128, 64, 32, 64, 32), + ], + "medium": [ + h_traits('fp32', 'fp32', 'fp32', 'fp32', 128, 128, 128, 128, 32, 128, 32), + # h_traits('fp32', 'fp32', 'fp32', 'fp32', 256, 128, 256, 128, 32, 128, 32), + ], + "large": [ + h_traits('fp32', 'fp32', 'fp32', 'fp32', 256, 128, 128, 128, 32, 128, 32), + # h_traits('fp32', 'fp32', 'fp32', 'fp32', 512, 128, 256, 256, 32, 256, 32), + ] + } + + # Toy example only support fp16 + dtype_combinations = [ + "fp16,fp16,fp16,fp16" + # "bf16,bf16,bf16,bf16" + ] + + total_blob = [] + for dtype_pair in dtype_combinations: + for size_category in trait_dict: + traits = trait_dict[size_category] + # Convert data types for the current dtype_pair + q_type, k_type, v_type, o_type = dtype_pair.split(',') + current_traits = [] + for t in traits: + new_t = copy.copy(t) + new_t.F_SaccDataType = 'fp32' # accumulation in fp32 + new_t.F_SMPLComputeDataType = 'fp32' # softmax compute in fp32 + new_t.F_PDataType = q_type + new_t.F_OaccDataType = 'fp32' # output accumulation in fp32 + current_traits.append(new_t) + + total_blob.append(h_instance(dtype_pair, size_category, current_traits)) + + return total_blob + + def list_blobs(self, args) -> None: + w_p = Path(self.working_path) + list_p = w_p / 'flash_attention_fwd_blobs.txt' + blobs = self.get_blobs(args) + + with list_p.open('w') as list_f: + # API related files + list_f.write(str(w_p / (self.name_api + ".cpp")) + "\n") + list_f.write(str(w_p / (self.name_common_header + ".hpp")) + "\n") + # Kernel instance files + for b in blobs: + list_f.write(str(w_p / (b.name + ".cpp")) + "\n") + + def gen_blobs(self, args) -> None: + w_p = Path(self.working_path) + w_str = self.content_api(args) + (w_p / (self.name_api + ".cpp")).write_text(w_str) + (w_p / (self.name_common_header + ".hpp")).write_text(self.content_common_header) + + blobs = self.get_blobs(args) + for b in blobs: + (w_p / (b.name + ".cpp")).write_text(b.content) + +def list_blobs(args): + api_list = args.api.split(',') + for api in api_list: + if api == 'fwd': + FlashAttentionFwdCodegen(args.working_path, args.filter).list_blobs(args) + +def gen_blobs(args): + api_list = args.api.split(',') + for api in api_list: + if api == 'fwd': + FlashAttentionFwdCodegen(args.working_path, args.filter).gen_blobs(args) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="generate", + description="gen API for Flash Attention kernel", + ) + parser.add_argument( + "-a", + "--api", + default='fwd', + required=False, + help="supply API(s) to generate (default: fwd). separated by comma." + ) + parser.add_argument( + "-w", + "--working_path", + default="./", + required=False, + help="the path where all the blobs are going to be generated" + ) + parser.add_argument( + "-l", + "--list_blobs", + action='store_true', + help="list all the kernels to a file" + ) + parser.add_argument( + "-g", + "--gen_blobs", + action='store_true', + help="generate all kernels into different tile" + ) + parser.add_argument( + "-f", + "--filter", + required=False, + help="filter out kernels that need to generate" + ) + + args = parser.parse_args() + + if (args.gen_blobs and args.list_blobs) or ((not args.gen_blobs) and (not args.list_blobs)): + print('gen_blobs/list_blobs must specify only one option') + sys.exit() + + p = Path(args.working_path) + if not p.exists(): + p.mkdir() + + if args.list_blobs: + list_blobs(args) + else: + gen_blobs(args)