mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-17 17:19:12 +00:00
Fix error after clang-format
This commit is contained in:
BoboFang
@@ -1,23 +1,22 @@
|
||||
set(EXAMPLE_REDUCE "add")
|
||||
#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}")
|
||||
# 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 add.cpp) target_include_directories(${
|
||||
EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR}) set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
|
||||
add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL add.cpp)
|
||||
target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
|
||||
set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
|
||||
|
||||
#generate assembly
|
||||
#list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS - v-- save - temps - Wno - gnu - line - marker)
|
||||
# generate assembly
|
||||
# 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)
|
||||
# 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)
|
||||
# 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)
|
||||
|
||||
@@ -1,28 +1,27 @@
|
||||
set(EXAMPLE_REDUCE "basic_gemm")
|
||||
#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}")
|
||||
# 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 gemm.cpp) target_include_directories(${
|
||||
EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR}) set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
|
||||
add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL gemm.cpp)
|
||||
target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${CMAKE_CURRENT_LIST_DIR})
|
||||
set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
|
||||
|
||||
#generate assembly
|
||||
#list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS - v-- save - temps - Wno - gnu - line - marker)
|
||||
# generate assembly
|
||||
# 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)
|
||||
# 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)
|
||||
|
||||
if(DEFINED kernel) message("Compiling with Kernel: ${kernel}")
|
||||
target_compile_definitions(${EXAMPLE_REDUCE} PRIVATE KERNEL_${kernel} = 1)
|
||||
endif()
|
||||
if(DEFINED kernel)
|
||||
message("Compiling with Kernel: ${kernel}")
|
||||
target_compile_definitions(${EXAMPLE_REDUCE} PRIVATE KERNEL_${kernel}=1)
|
||||
endif()
|
||||
|
||||
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)
|
||||
# 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)
|
||||
|
||||
31
example/ck_tile/99_toy_example/03_flash_attention_fwd/CMakeLists.txt
Normal file → Executable file
31
example/ck_tile/99_toy_example/03_flash_attention_fwd/CMakeLists.txt
Normal file → Executable file
@@ -1,22 +1,19 @@
|
||||
set(EXAMPLE_REDUCE "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}")
|
||||
# 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})
|
||||
set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
|
||||
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)
|
||||
|
||||
#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)
|
||||
# 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)
|
||||
# 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)
|
||||
|
||||
85
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
Normal file → Executable file
85
example/ck_tile/99_toy_example/04_codegen_flash_attention_fwd/CMakeLists.txt
Normal file → Executable file
@@ -1,48 +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()
|
||||
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()
|
||||
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)
|
||||
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)
|
||||
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") message(
|
||||
"adding example ${EXAMPLE_REDUCE}")
|
||||
set(EXAMPLE_REDUCE "codegen_basic_flash_attention_fwd")
|
||||
message("adding example ${EXAMPLE_REDUCE}")
|
||||
|
||||
add_executable(${EXAMPLE_REDUCE} EXCLUDE_FROM_ALL flash_attention_fwd.cpp)
|
||||
add_executable(${EXAMPLE_REDUCE}
|
||||
EXCLUDE_FROM_ALL
|
||||
flash_attention_fwd.cpp
|
||||
)
|
||||
|
||||
target_include_directories(${EXAMPLE_REDUCE} PRIVATE ${
|
||||
CMAKE_CURRENT_LIST_DIR})
|
||||
target_include_directories(${EXAMPLE_REDUCE}
|
||||
PRIVATE
|
||||
${CMAKE_CURRENT_LIST_DIR}
|
||||
)
|
||||
|
||||
target_sources(${EXAMPLE_REDUCE} PRIVATE ${
|
||||
FLASH_ATTENTION_FWD_GEN_BLOBS})
|
||||
target_sources(${EXAMPLE_REDUCE} PRIVATE ${FLASH_ATTENTION_FWD_GEN_BLOBS})
|
||||
|
||||
message("FLASH_ATTENTION_FWD_GEN_BLOBS = "
|
||||
"${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)
|
||||
|
||||
target_compile_options(${EXAMPLE_REDUCE} PRIVATE ${
|
||||
EXAMPLE_REDUCE_COMPILE_OPTIONS})
|
||||
set(EXAMPLE_REDUCE_COMPILE_OPTIONS)
|
||||
list(APPEND EXAMPLE_REDUCE_COMPILE_OPTIONS
|
||||
-Wno-undefined-func-template
|
||||
-Wno-float-equal
|
||||
--offload-compress
|
||||
)
|
||||
|
||||
set_property(GLOBAL PROPERTY RULE_MESSAGES OFF)
|
||||
target_compile_options(${EXAMPLE_REDUCE}
|
||||
PRIVATE
|
||||
${EXAMPLE_REDUCE_COMPILE_OPTIONS}
|
||||
)
|
||||
|
||||
set_property(GLOBAL PROPERTY RULE_MESSAGES OFF)
|
||||
@@ -1,15 +1,23 @@
|
||||
#SPDX - License - Identifier : MIT
|
||||
#Copyright(c) 2025, Advanced Micro Devices, Inc.All rights reserved.
|
||||
# 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
|
||||
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(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":
|
||||
@@ -26,7 +34,7 @@ def BOOL_MAP(b_) -> str:
|
||||
|
||||
class FlashAttentionFwdCodegen:
|
||||
API_TRAITS_DEFINE = """
|
||||
|
||||
|
||||
template <typename SaccDataType_,
|
||||
typename SMPLComputeDataType_,
|
||||
typename PDataType_,
|
||||
@@ -52,12 +60,12 @@ struct flash_attention_fwd_traits_
|
||||
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 <typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
@@ -82,80 +90,77 @@ using traits_ = flash_attention_fwd_traits_<SaccDataType,
|
||||
kK1PerBlock>;
|
||||
"""
|
||||
|
||||
#API_COMMON_HEADER = "" \
|
||||
"
|
||||
#// SPDX-License-Identifier: MIT
|
||||
#// Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
|
||||
# API_COMMON_HEADER = """
|
||||
# // SPDX-License-Identifier: MIT
|
||||
# // Copyright (c) 2024, Advanced Micro Devices, Inc. All rights reserved.
|
||||
|
||||
# #include<ck_tile / core.hpp>
|
||||
# #include <ck_tile/core.hpp>
|
||||
# #include "flash_attention_fwd.hpp"
|
||||
# #include<iostream>
|
||||
# #include <iostream>
|
||||
|
||||
# #pragma once
|
||||
|
||||
#using S = ck_tile::stream_config;
|
||||
#using A = FlashAttnArgs;
|
||||
# using S = ck_tile::stream_config;
|
||||
# using A = FlashAttnArgs;
|
||||
|
||||
#{F_traits_define }
|
||||
# {F_traits_define}
|
||||
|
||||
#template < typename QDataType,
|
||||
#typename KDataType,
|
||||
#typename VDataType,
|
||||
#typename ODataType,
|
||||
#typename Traits_>
|
||||
#float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType, ODataType>& 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;
|
||||
# template <typename QDataType,
|
||||
# typename KDataType,
|
||||
# typename VDataType,
|
||||
# typename ODataType,
|
||||
# typename Traits_>
|
||||
# float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType, ODataType>& 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);
|
||||
|
||||
#if (stream_config.log_level_ > 0)
|
||||
#std::cout << ", " << "FlashAttentionFwd<" << Traits_::kBlockSize << "," << Traits_::kHeadDim \
|
||||
<< ">" << std::flush;
|
||||
# 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 < Traits_::kBlockSize, Traits_::kBlockPerCu>(
|
||||
#ck_tile::FlashAttentionFwd < QDataType,
|
||||
#KDataType,
|
||||
#VDataType,
|
||||
#SaccDataType,
|
||||
#SMPLComputeDataType,
|
||||
#PDataType,
|
||||
#OaccDataType,
|
||||
#ODataType,
|
||||
#Traits_::kBlockSize,
|
||||
#Traits_::kHeadDim,
|
||||
#Traits_::kM0PerBlock,
|
||||
#Traits_::kN0PerBlock,
|
||||
#Traits_::kK0PerBlock,
|
||||
#Traits_::kN1PerBlock,
|
||||
#Traits_::kK1PerBlock>{{} },
|
||||
#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
|
||||
#} }
|
||||
#"" \
|
||||
"
|
||||
# return ck_tile::launch_kernel(stream_config,
|
||||
# ck_tile::make_kernel<Traits_::kBlockSize, Traits_::kBlockPerCu>(
|
||||
# ck_tile::FlashAttentionFwd<QDataType,
|
||||
# KDataType,
|
||||
# VDataType,
|
||||
# SaccDataType,
|
||||
# SMPLComputeDataType,
|
||||
# PDataType,
|
||||
# OaccDataType,
|
||||
# ODataType,
|
||||
# Traits_::kBlockSize,
|
||||
# Traits_::kHeadDim,
|
||||
# Traits_::kM0PerBlock,
|
||||
# Traits_::kN0PerBlock,
|
||||
# Traits_::kK0PerBlock,
|
||||
# Traits_::kN1PerBlock,
|
||||
# Traits_::kK1PerBlock>{{}},
|
||||
# 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
|
||||
@@ -199,19 +204,14 @@ template float flash_attention_fwd<ck_tile::half_t, ck_tile::half_t, ck_tile::ha
|
||||
}}
|
||||
"""
|
||||
|
||||
#API_PER_DTYPE = \
|
||||
"" \
|
||||
" {F_if}(std::is_same_v<QDataType, {F_q_type}> && std::is_same_v<KDataType, {F_k_type}> && std::is_same_v<VDataType, {F_v_type}> && std::is_same_v<ODataType, {F_o_type}>) {{
|
||||
#{F_per_size_case }
|
||||
#} }
|
||||
#"" \
|
||||
"
|
||||
#API_PER_SIZE_CASE = "" \
|
||||
" {F_if} {F_SIZE_COND} {{
|
||||
#{F_inner_dispatch }
|
||||
#} }
|
||||
#"" \
|
||||
"
|
||||
# API_PER_DTYPE = """ {F_if}(std::is_same_v<QDataType, {F_q_type}> && std::is_same_v<KDataType, {F_k_type}> && std::is_same_v<VDataType, {F_v_type}> && std::is_same_v<ODataType, {F_o_type}>) {{
|
||||
# {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_<QDataType, KDataType, VDataType, ODataType, traits_<{F_trait_name}>>(a, stream_config);
|
||||
"""
|
||||
@@ -224,7 +224,7 @@ template float flash_attention_fwd<ck_tile::half_t, ck_tile::half_t, ck_tile::ha
|
||||
|
||||
namespace ck_tile {
|
||||
// clang-format off
|
||||
//
|
||||
//
|
||||
{F_instance_def}
|
||||
// clang-format on
|
||||
|
||||
@@ -315,7 +315,7 @@ namespace ck_tile {{
|
||||
#include "flash_attention_fwd.hpp"
|
||||
|
||||
namespace ck_tile {{
|
||||
|
||||
|
||||
template <typename SaccDataType_,
|
||||
typename SMPLComputeDataType_,
|
||||
typename PDataType_,
|
||||
@@ -341,13 +341,13 @@ struct flash_attention_fwd_traits_
|
||||
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 <typename SaccDataType,
|
||||
typename SMPLComputeDataType,
|
||||
typename PDataType,
|
||||
@@ -379,11 +379,11 @@ template <typename QDataType,
|
||||
typename Traits_>
|
||||
float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType, ODataType>& 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;
|
||||
|
||||
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)
|
||||
@@ -430,10 +430,10 @@ float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType,
|
||||
}}
|
||||
"""
|
||||
def content_api(self, args) -> str:
|
||||
#Sort based on dtype
|
||||
# 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] = {}
|
||||
@@ -445,16 +445,16 @@ float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType,
|
||||
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)"
|
||||
@@ -462,30 +462,30 @@ float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType,
|
||||
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(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 += 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
|
||||
#)
|
||||
# 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(
|
||||
@@ -498,33 +498,33 @@ float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType,
|
||||
h_traits = self.h_traits
|
||||
h_instance = self.h_instance
|
||||
|
||||
#Define kernel configurations for different size categories
|
||||
# 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),
|
||||
# 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),
|
||||
# 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),
|
||||
# h_traits('fp32', 'fp32', 'fp32', 'fp32', 512, 128, 256, 256, 32, 256, 32),
|
||||
]
|
||||
}
|
||||
|
||||
#Toy example only support fp16
|
||||
# Toy example only support fp16
|
||||
dtype_combinations = [
|
||||
"fp16,fp16,fp16,fp16"
|
||||
#"bf16,bf16,bf16,bf16"
|
||||
# "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
|
||||
# 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:
|
||||
@@ -534,21 +534,21 @@ float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType,
|
||||
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
|
||||
# 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
|
||||
# Kernel instance files
|
||||
for b in blobs:
|
||||
list_f.write(str(w_p / (b.name + ".cpp")) + "\n")
|
||||
|
||||
@@ -557,7 +557,7 @@ float flash_attention_fwd_(const FlashAttnArgs<QDataType, KDataType, VDataType,
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user