saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

ROCm Flash Attention (#14838)

Browse source 
Adds flash attention via composable kernel for ROCm EP
This commit is contained in:
cloudhan 2023-03-16 10:39:58 +08:00 committed by GitHub
parent 881f3f6be3
commit a5ab88247b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
20 changed files with 1304 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cmake/external/composable_kernel.cmake vendored
View file

@ -1,5 +1,5 @@
set(composable_kernel_URL https://github.com/ROCmSoftwarePlatform/composable_kernel.git)
set(composable_kernel_TAG 0345963eef4f92e9c5eab608bb8557b5463a1dcb) # 2022-12-15 15:07:24 -0600
set(composable_kernel_TAG bef0cb20dba0d9b315df46899310478a81c21852) # 2023-02-16 11:54:08 -0800
set(PATCH ${PROJECT_SOURCE_DIR}/patches/composable_kernel/Fix_Clang_Build.patch)

5
cmake/onnxruntime_providers.cmake
View file

@ -1508,8 +1508,9 @@ if (onnxruntime_USE_ROCM)
target_link_libraries(onnxruntime_providers_rocm PRIVATE
onnxruntime_composable_kernel_includes
# Currently we shall not use composablekernels::device_operations, the target includes all conv dependencies, which
# are extremely slow to compile. Instead, we only link all gemm related objects. See the following link on updating.
# https://github.com/ROCmSoftwarePlatform/composable_kernel/blob/85978e0201/library/src/tensor_operation_instance/gpu/CMakeLists.txt#L33-L54
# are extremely slow to compile. Instead, we only link all gemm related objects. See the following directory on
# updating.
# https://github.com/ROCmSoftwarePlatform/composable_kernel/tree/develop/library/src/tensor_operation_instance/gpu
device_gemm_instance
device_gemm_add_fastgelu_instance
device_gemm_fastgelu_instance

10
cmake/patches/composable_kernel/Fix_Clang_Build.patch
View file

@ -1,5 +1,5 @@
diff --git a/CMakeLists.txt b/CMakeLists.txt
index f861e302..f0b6bcea 100644
index f861e3020..f0b6bceae 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,7 +1,7 @@
@ -71,14 +71,18 @@ index f861e302..f0b6bcea 100644
- HEADER_ONLY
-)
diff --git a/library/src/tensor_operation_instance/gpu/CMakeLists.txt b/library/src/tensor_operation_instance/gpu/CMakeLists.txt
index c206c4dc..e45fac9d 100644
index c206c4dc0..b283eeb64 100644
--- a/library/src/tensor_operation_instance/gpu/CMakeLists.txt
+++ b/library/src/tensor_operation_instance/gpu/CMakeLists.txt
@@ -1,7 +1,9 @@
@@ -1,7 +1,13 @@
function(add_instance_library INSTANCE_NAME)
message("adding instance ${INSTANCE_NAME}")
+ set_source_files_properties(${ARGN} PROPERTIES LANGUAGE HIP)
add_library(${INSTANCE_NAME} OBJECT ${ARGN})
+ # Always disable debug symbol and C debug assert due to
+ # - Linker error: ... relocation truncated to fit ..., caused by object files to be linked are too huge.
+ # - https://github.com/ROCmSoftwarePlatform/composable_kernel/issues/622
+ target_compile_options(${INSTANCE_NAME} PRIVATE -g0 -DNDEBUG)
target_compile_features(${INSTANCE_NAME} PUBLIC)
+ target_compile_definitions(${INSTANCE_NAME} PRIVATE "__HIP_PLATFORM_AMD__=1" "__HIP_PLATFORM_HCC__=1")
set_target_properties(${INSTANCE_NAME} PROPERTIES POSITION_INDEPENDENT_CODE ON)

41
onnxruntime/contrib_ops/rocm/bert/attention.cu
View file

@ -83,13 +83,17 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
using HipT = typename ToHipType<T>::MappedType;
using QkvProjectGeneric = GemmPermuteGenericPipeline<HipT>;
using AttentionGeneric = GemmSoftmaxGemmPermuteGenericPipeline<HipT>;
using AttentionTunableOp = GemmSoftmaxGemmPermuteTunableOp<HipT>;
size_t qkv_project_output_bytes = QkvProjectGeneric::GetOutputNumBytes(&attn);
size_t attention_workspace_bytes = AttentionGeneric::GetWorkspaceNumBytes(&attn);
ORT_ENFORCE(QkvProjectGeneric::GetWorkspaceNumBytes(&attn) <= attention_workspace_bytes); // workspace reuse
size_t shared_workspace_bytes = std::max(QkvProjectGeneric::GetWorkspaceNumBytes(&attn),
AttentionGeneric::GetWorkspaceNumBytes(&attn));
if (GetTuningContext()->IsTunableOpEnabled()) {
shared_workspace_bytes = std::max(shared_workspace_bytes, AttentionTunableOp::GetWorkspaceNumBytes(&attn));
}
auto qkv_project_output = GetScratchBuffer<void>(qkv_project_output_bytes, context->GetComputeStream());
auto workspace = GetScratchBuffer<void>(attention_workspace_bytes, context->GetComputeStream());
auto workspace = GetScratchBuffer<void>(shared_workspace_bytes, context->GetComputeStream());
GemmPermuteParams<HipT> gemm_permute_params;
{
@ -105,7 +109,7 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
params.bias_buffer = reinterpret_cast<const HipT*>(bias->DataRaw());
params.out_buffer = reinterpret_cast<HipT*>(qkv_project_output.get());
params.ones = GetConstOnes<HipT>(attn.batch_size * attn.sequence_length, stream);
params.workspace_buffer = reinterpret_cast<HipT*>(workspace.get()); // workspace reuse
params.workspace_buffer = reinterpret_cast<HipT*>(workspace.get());
}
ORT_RETURN_IF_ERROR(QkvProjectGeneric::Run(&gemm_permute_params));
@ -118,16 +122,16 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
const int batches = attn.batch_size * attn.num_heads;
const int present_size_per_batch = attn.total_sequence_length * attn.head_size;
ORT_RETURN_IF_ERROR(
LaunchConcatPastToPresent(Stream(context),
attn.total_sequence_length,
attn.sequence_length,
attn.batch_size,
attn.head_size,
attn.num_heads,
device_prop.maxThreadsPerBlock,
nullptr == past ? nullptr : reinterpret_cast<const HipT*>(past->DataRaw()),
k_buffer,
reinterpret_cast<HipT*>(present->MutableDataRaw())));
LaunchConcatPastToPresent(Stream(context),
attn.total_sequence_length,
attn.sequence_length,
attn.batch_size,
attn.head_size,
attn.num_heads,
device_prop.maxThreadsPerBlock,
nullptr == past ? nullptr : reinterpret_cast<const HipT*>(past->DataRaw()),
k_buffer,
reinterpret_cast<HipT*>(present->MutableDataRaw())));
// update pointers to present_k and present_v.
k_buffer = reinterpret_cast<HipT*>(present->MutableDataRaw());
@ -159,13 +163,18 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
if (mask_index != nullptr) {
params.mask_index_buffer = mask_index->Data<int>();
params.mask_index_dims = mask_index->Shape().GetDims();
params.mask_index_dims = mask_index->Shape().AsShapeVector();
}
params.workspace_buffer = reinterpret_cast<HipT*>(workspace.get());
}
return AttentionGeneric::Run(&gemm_softmax_gemm_permute_params, use_persistent_softmax);
if (this->GetTuningContext()->IsTunableOpEnabled() &&
!use_persistent_softmax) {
return AttentionTunableOp{}(&gemm_softmax_gemm_permute_params);
} else {
return AttentionGeneric::Run(&gemm_softmax_gemm_permute_params, use_persistent_softmax);
}
}
} // namespace rocm

4
onnxruntime/contrib_ops/rocm/bert/attention_impl.cu
View file

@ -67,10 +67,10 @@ size_t GetAttentionWorkspaceSize(
int num_heads,
int head_size,
int sequence_length,
int past_sequence_length) {
int total_sequence_length) {
size_t qkv_size = element_size * 3 * batch_size * sequence_length * num_heads * head_size;
return qkv_size + 2 * GetAttentionScratchSize(element_size, batch_size, num_heads,
sequence_length, past_sequence_length + sequence_length);
sequence_length, total_sequence_length);
}
inline int3 Get2DMaskStrides(int total_sequence_length) {

17
onnxruntime/contrib_ops/rocm/bert/attention_softmax.h
View file

@ -38,7 +38,6 @@ namespace rocm {
template <typename T, unsigned TPB>
__device__ inline void Softmax(const int all_sequence_length,
const int sequence_length,
const int valid_end,
const int valid_start,
const T* add_before_softmax,
@ -276,10 +275,8 @@ __global__ void SoftmaxKernelSmall(const int all_sequence_length, const int sequ
}
template <typename T, unsigned TPB>
__global__ void SoftmaxKernel(const int all_sequence_length, const int sequence_length,
const T* add_before_softmax, const T* input, T* output) {
Softmax<T, TPB>(all_sequence_length, sequence_length, all_sequence_length, 0,
add_before_softmax, input, output);
__global__ void SoftmaxKernel(const int all_sequence_length, const T* add_before_softmax, const T* input, T* output) {
Softmax<T, TPB>(all_sequence_length, all_sequence_length, 0, add_before_softmax, input, output);
}
template <typename T>
@ -315,7 +312,7 @@ Status ComputeSoftmax(
} else if (!is_unidirectional) {
const int blockSize = 1024;
SoftmaxKernel<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, add_before_softmax, input, output);
all_sequence_length, add_before_softmax, input, output);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention ROCM operator does not support total sequence length > 1024.");
}
@ -349,8 +346,7 @@ __global__ void MaskedSoftmaxKernelSmall(const int all_sequence_length, const in
}
template <typename T, unsigned TPB>
__global__ void MaskedSoftmaxKernel(const int all_sequence_length, const int sequence_length,
const int* mask_end, const int* mask_start,
__global__ void MaskedSoftmaxKernel(const int all_sequence_length, const int* mask_end, const int* mask_start,
const T* add_before_softmax, const T* input, T* output) {
__shared__ int start_position;
__shared__ int end_position;
@ -368,8 +364,7 @@ __global__ void MaskedSoftmaxKernel(const int all_sequence_length, const int seq
}
__syncthreads();
Softmax<T, TPB>(all_sequence_length, sequence_length, end_position, start_position,
add_before_softmax, input, output);
Softmax<T, TPB>(all_sequence_length, end_position, start_position, add_before_softmax, input, output);
}
template <typename T>
@ -400,7 +395,7 @@ Status ComputeSoftmaxWithMask1D(
} else if (!is_unidirectional) {
const int blockSize = 1024;
MaskedSoftmaxKernel<T, blockSize><<<grid, blockSize, 0, stream>>>(
all_sequence_length, sequence_length, mask_index, mask_start,
all_sequence_length, mask_index, mask_start,
add_before_softmax, input, output);
} else {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Attention ROCM operator does not support total sequence length > 1024.");

146
onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh Normal file
View file

@ -0,0 +1,146 @@
// SPDX-License-Identifier: MIT
// Modifications Copyright (c) Microsoft Corporation. All rights reserved.
// Copyright (c) 2018-2022, Advanced Micro Devices, Inc. All rights reserved.
#pragma once
#include <cstdlib>
#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
#include "ck/tensor_operation/gpu/device/tensor_specialization.hpp"
#include "ck/tensor_operation/gpu/device/device_batched_gemm_softmax_gemm_permute.hpp"
#include "ck/tensor_operation/gpu/device/gemm_specialization.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp"
#include "ck/tensor_operation/gpu/element/unary_element_wise_operation.hpp"
#include "ck/utility/data_type.hpp"
namespace onnxruntime {
namespace contrib {
namespace rocm {
namespace internal {
using F16 = ck::half_t;
using F32 = float;
template <ck::index_t... Is>
using S = ck::Sequence<Is...>;
using MaskingSpecialization = ck::tensor_operation::device::MaskingSpecialization;
using PassThrough = ck::tensor_operation::element_wise::PassThrough;
using ck::tensor_operation::device::DeviceBatchedGemmSoftmaxGemmPermute; // the interface
using ck::tensor_operation::device::DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle; // the implementation
static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
static constexpr auto GemmPadded = ck::tensor_operation::device::GemmSpecialization::MNKOPadding;
static constexpr auto TensorDefault = ck::tensor_operation::device::TensorSpecialization::Default;
template <ck::index_t NumDimG,
ck::index_t NumDimM,
ck::index_t NumDimN,
ck::index_t NumDimK,
ck::index_t NumDimO,
typename DT, // A, B0, B1, C, CShuffle Datatype
typename D0sDT, // D0 (bias in A*B0+D0s) DataType
typename AccDT, // Accumulator Datatype
typename D0Op, // D0 DataType
MaskingSpecialization MaskingSpec>
using device_batched_gemm_softmax_gemm_permute_instances =
std::tuple<
// clang-format off
// #############################################| NumDimG| NumDimM| NumDimN| NumDimK| NumDimO| AData| B0Data| B1Data| CData| Acc0BiasData| Acc1BiasData| AccData| CShuffle| A| B0| Acc0| B1| C| GEMM| ATensorSpec| B0TensorSpec| B1TensorSpec| CTensorSpec| NumGemmK| Block| Gemm01| Gemm0| Gemm0| Gemm1| Gemm1| AK1| BK1| B1K1| MPer| NPer| Gemm0| Gemm0| Gemm1| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockTransfer| ABlockLds| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockTransfer| B0BlockLds| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockTransfer| B1BlockLds| CShuffle| CShuffle| CBlockTransferClusterLengths| CBlockTransfer| MaskingSpec|
// #############################################| | | | | | Type| Type| Type| Type| Type| Type| Type| DataType| Elementwise| Elementwise| Elementwise| Elementwise| Elementwise| Specialization| | | | | Prefetch| Size| MPer| NPer| KPer| NPer| KPer| | | | XDL| XDL| MXdl| NXdl| NXdl| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraM| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| ThreadCluster| ThreadCluster| SrcAccessOrder| SrcVectorDim| SrcScalar| DstScalar| AddExtraN| MXdlPerWave| NXdlPerWave| _MBlock_MWaveMPerXdl| ScalarPerVector| |
// #############################################| | | | | | | | | | | | | | Operation| Operation| Operation| Operation| Operation| | | | | | Stage| | Block| Block| Block| Block| Block| | | | | | Per| Per| Per| Lengths_K0_M_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | Lengths_K0_N_K1| ArrangeOrder| | | PerVector| PerVector_K1| | PerShuffle| PerShuffle| _NBlock_NWaveNPerXdl| _NWaveNPerXdl| |
// #############################################| | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | | Wave| Wave| Wave| | | | | | | | | | | | | | | | | | | | | | | | | | |
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 256, 128, 32, 64, 32, 8, 8, 2, 32, 32, 2, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 256, 128, 32, 128, 32, 8, 8, 2, 32, 32, 2, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
#if ROCM_VERSION >= 50500
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 256, 32, 64, 32, 8, 8, 2, 32, 32, 1, 8, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
#endif
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 256, 32, 128, 32, 8, 8, 2, 32, 32, 1, 8, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 128, 64, 64, 32, 8, 8, 2, 32, 32, 1, 4, 2, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 128, 32, 64, 32, 8, 8, 2, 32, 32, 1, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 128, 64, 128, 32, 8, 8, 2, 32, 32, 1, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 128, 32, 128, 32, 8, 8, 2, 32, 32, 1, 4, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 64, 256, 32, 128, 32, 8, 8, 2, 16, 16, 1, 16, 8, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 8, S<1, 16, 1,16>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 64, 256, 32, 64, 32, 8, 8, 2, 16, 16, 1, 16, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 4, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 64, 256, 64, 128, 32, 8, 8, 2, 16, 16, 1, 16, 8, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 8, S<1, 16, 1,16>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmDefault, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 64, 256, 64, 64, 32, 8, 8, 2, 16, 16, 1, 16, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 4, S<1, 32, 1, 8>, 8, MaskingSpec>,
// Padded fallback kernel
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 128, 64, 128, 32, 8, 8, 2, 32, 32, 1, 4, 4, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S<8, 32, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, false, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>,
DeviceBatchedGemmSoftmaxGemmPermute_Xdl_CShuffle< NumDimG, NumDimM, NumDimN, NumDimK, NumDimO, DT, DT, DT, DT, D0sDT, ck::Tuple<>, AccDT, DT, PassThrough, PassThrough, D0Op, PassThrough, PassThrough, GemmPadded, TensorDefault, TensorDefault, TensorDefault, TensorDefault, 1, 256, 128, 64, 32, 128, 32, 8, 8, 2, 32, 32, 1, 2, 4, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S< 8, 32, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, MaskingSpec>
// clang-format on
>;
struct PreSoftmaxAttentionScoreOp {
PreSoftmaxAttentionScoreOp(float scale) : scale_(scale) {}
// non-biased, non-masked
__host__ __device__ void operator()(float& y, const float& x) const {
y = scale_ * x;
}
// biased or converted masked
__host__ __device__ void operator()(float& y, const float& x, const F16& bias) const {
y = scale_ * x + ck::type_convert<float>(bias);
}
// biased and converted masked
__host__ __device__ void operator()(float& y, const float& x, const F16& bias, const F16& converted_mask) const {
y = scale_ * x + ck::type_convert<float>(bias) + ck::type_convert<float>(converted_mask);
}
float scale_;
};
// Use this function to gat implementation
template <typename DT, typename D0sDT, typename AccDT, typename D0Op, MaskingSpecialization MaskingSpec>
std::vector<std::unique_ptr<DeviceBatchedGemmSoftmaxGemmPermute<
2, 1, 1, 1, 1,
DT, DT, DT, DT, D0sDT, ck::Tuple<>,
PassThrough, PassThrough, D0Op, PassThrough, PassThrough,
MaskingSpec>>>
GetDeviceBatchedGemmSoftmaxGemmPermuteInstances() {
return {};
}
// implemented in impl_{fp16,bf16}[_biased][_masked].cu
// fp16, non-biased, non-masked
template <>
std::vector<std::unique_ptr<DeviceBatchedGemmSoftmaxGemmPermute<
2, 1, 1, 1, 1,
F16, F16, F16, F16, ck::Tuple<>, ck::Tuple<>,
PassThrough, PassThrough, PreSoftmaxAttentionScoreOp, PassThrough, PassThrough,
MaskingSpecialization::MaskDisabled>>>
GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
F16, ck::Tuple<>, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>();
// fp16, biased, non-masked
template <>
std::vector<std::unique_ptr<DeviceBatchedGemmSoftmaxGemmPermute<
2, 1, 1, 1, 1,
F16, F16, F16, F16, ck::Tuple<F16>, ck::Tuple<>,
PassThrough, PassThrough, PreSoftmaxAttentionScoreOp, PassThrough, PassThrough,
MaskingSpecialization::MaskDisabled>>>
GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
F16, ck::Tuple<F16>, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>();
// fp16, biased, fp16 masked, basically, two bias
template <>
std::vector<std::unique_ptr<DeviceBatchedGemmSoftmaxGemmPermute<
2, 1, 1, 1, 1,
F16, F16, F16, F16, ck::Tuple<F16, F16>, ck::Tuple<>,
PassThrough, PassThrough, PreSoftmaxAttentionScoreOp, PassThrough, PassThrough,
MaskingSpecialization::MaskDisabled>>>
GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
F16, ck::Tuple<F16, F16>, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>();
} // namespace ck
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

37
onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16.cu Normal file
View file

@ -0,0 +1,37 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp"
namespace onnxruntime {
namespace contrib {
namespace rocm {
namespace internal {
using NonBiasedNonmasked = DeviceBatchedGemmSoftmaxGemmPermute<
2, 1, 1, 1, 1,
F16, F16, F16, F16, ck::Tuple<>, ck::Tuple<>,
PassThrough, PassThrough, PreSoftmaxAttentionScoreOp, PassThrough, PassThrough,
MaskingSpecialization::MaskDisabled>;
template <>
std::vector<std::unique_ptr<NonBiasedNonmasked>>
GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
F16, ck::Tuple<>, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>() {
std::vector<std::unique_ptr<NonBiasedNonmasked>> instances;
ck::tensor_operation::device::instance::add_device_operation_instances(
instances,
device_batched_gemm_softmax_gemm_permute_instances<
2, 1, 1, 1, 1,
F16, ck::Tuple<>, F32, PreSoftmaxAttentionScoreOp,
MaskingSpecialization::MaskDisabled>{});
return instances;
}
} // namespace internal
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

37
onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased.cu Normal file
View file

@ -0,0 +1,37 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp"
namespace onnxruntime {
namespace contrib {
namespace rocm {
namespace internal {
using BiasedNonmasked = DeviceBatchedGemmSoftmaxGemmPermute<
2, 1, 1, 1, 1,
F16, F16, F16, F16, ck::Tuple<F16>, ck::Tuple<>,
PassThrough, PassThrough, PreSoftmaxAttentionScoreOp, PassThrough, PassThrough,
MaskingSpecialization::MaskDisabled>;
template <>
std::vector<std::unique_ptr<BiasedNonmasked>>
GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
F16, ck::Tuple<F16>, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>() {
std::vector<std::unique_ptr<BiasedNonmasked>> instances;
ck::tensor_operation::device::instance::add_device_operation_instances(
instances,
device_batched_gemm_softmax_gemm_permute_instances<
2, 1, 1, 1, 1,
F16, ck::Tuple<F16>, F32, PreSoftmaxAttentionScoreOp,
MaskingSpecialization::MaskDisabled>{});
return instances;
}
} // namespace internal
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

37
onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl_fp16_biased_biased.cu Normal file
View file

@ -0,0 +1,37 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
#include "ck/library/tensor_operation_instance/add_device_operation_instance.hpp"
#include "ck/tensor_operation/gpu/device/impl/device_batched_gemm_softmax_gemm_permute_xdl_cshuffle.hpp"
namespace onnxruntime {
namespace contrib {
namespace rocm {
namespace internal {
using BiasedNonmasked = DeviceBatchedGemmSoftmaxGemmPermute<
2, 1, 1, 1, 1,
F16, F16, F16, F16, ck::Tuple<F16, F16>, ck::Tuple<>,
PassThrough, PassThrough, PreSoftmaxAttentionScoreOp, PassThrough, PassThrough,
MaskingSpecialization::MaskDisabled>;
template <>
std::vector<std::unique_ptr<BiasedNonmasked>>
GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
F16, ck::Tuple<F16, F16>, F32, PreSoftmaxAttentionScoreOp, MaskingSpecialization::MaskDisabled>() {
std::vector<std::unique_ptr<BiasedNonmasked>> instances;
ck::tensor_operation::device::instance::add_device_operation_instances(
instances,
device_batched_gemm_softmax_gemm_permute_instances<
2, 1, 1, 1, 1,
F16, ck::Tuple<F16, F16>, F32, PreSoftmaxAttentionScoreOp,
MaskingSpecialization::MaskDisabled>{});
return instances;
}
} // namespace internal
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

303
onnxruntime/contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh
View file

@ -18,10 +18,10 @@ In QKV projection (prior to this pipeline):
X --o--> K [B,T,N*H] ->Reshape-> [B,T,N,H] ->Permute0213-> [B,N,T,H]
\-> V [B,T,N*H] ->Reshape-> [B,T,N,H] ->Permute0213-> [B,N,T,H]
pre_softmax_attn_scores = Q*K' = [B,N,S,H] * [BxNxTxH]' = [B,N,S,T] Batched GEMM1
pre_softmax_attn_scores_masked = pre_softmax_attn_scores +? bias +? mask Add Bias, +? is optional
attn_scores = softmax(pre_softmax_attn_scores_masked * scale) = [B,N,S,T] Scale then Softmax
scaled_multi_head_attn = attn_scores * V = [B,N,S,T] * [B,N,T,H] = [B,N,S,H] Batched GEMM2
pre_softmax_attn_scores = Q*K' = [B,N,S,H] * [BxNxTxH]' = [B,N,S,T] Batched GEMM1
pre_softmax_attn_scores_masked = pre_softmax_attn_scores * scale +? bias +? mask Scale Add Bias, +? is optional
attn_scores = softmax(pre_softmax_attn_scores_masked) = [B,N,S,T] Softmax
scaled_multi_head_attn = attn_scores * V = [B,N,S,T] * [B,N,T,H] = [B,N,S,H] Batched GEMM2
Op outputs scaled_multi_head_attn:
[B,N,S,H] ->Permute0213-> [B,S,N,H] ->Reshape-> [B,S,N*H]
@ -31,13 +31,46 @@ For the computing of pre_softmax_attn_scores +? mask +? bias:
GemmSoftmaxGemmPermuteGenericPipeline handles it in specialized softmax. TODO: remove it!
CK in GemmSoftmaxGemmPermuteTunablePipeline
Q*K' ---> scale ---> [B,N,S,T] -------+?--> masked
bias --------------> [B,N,S,T] --+?--/
mask_2d ---> [B,T] ---> [B,1,1,T] -/
Q*K' ---> scale ---> [B,N,S,T] -------+?--> masked
bias --------------> [B,N,S,T] --+?--/
mask_3d --> [B,S,T] --> [B,1,S,T] -/
Q*K' ---> scale ---> [B,N,S,T] -------+?--> masked
bias --------------> [B,N,S,T] --+?--/
mask_4d -> [B,1,M,M] -> [B,1,S,T] -/ M is max_sequence_length from megatron, we will create a
**sub-view** from original mask buffer
For CK implementation, there will be four cases combined:
non-biased, non-masked, no special processing.
biased, non-masked, no special processing, add the mask directly.
non-biased, masked, convert the mask to [B,1,1_or_S,T] and perform broadcast add with scaled Q*K'.
biased, masked, convert the mask to [B,1,1_or_S,T] and perform broadcast add with bias and scaled Q*K'.
Broadcast add is not actually perform the broadcasting, just broadcast the load operation from memory. The impl details
are in composable kernels. The scale and add logic is performed via Acc0ElementOp
*/
#include "core/framework/tensor_shape.h"
#include "core/providers/rocm/tunable/gemm.h"
#include "core/providers/rocm/tunable/rocm_tunable.h"
#include "contrib_ops/cpu/bert/attention_base.h"
#include "contrib_ops/rocm/bert/attention_impl.h"
#include "contrib_ops/rocm/bert/attention_softmax.h"
#include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_ck_impl/impl.cuh"
#include "ck/ck.hpp"
#include "ck/tensor_operation/gpu/device/tensor_layout.hpp"
#include "ck/tensor_operation/gpu/element/element_wise_operation.hpp"
#include <array>
#include <vector>
namespace blas = onnxruntime::rocm::tunable::blas;
@ -53,7 +86,7 @@ inline int3 Get2DMaskStrides(int total_sequence_length) {
inline std::tuple<const int*, int3, int3> GetRawMaskBufferAddrSizesAndStrides(
const int* buffer, const AttentionParameters* attn) {
const int* offseted_buffer{buffer}; // how to view the mask buffer
int3 sizes{-1, -1, -1}; // the logical shape of the view
int3 sizes{0, 0, 0}; // the logical shape of the view
int3 strides{-1, -1, -1}; // the physical memory layout
switch (attn->mask_type) {
case MASK_NONE:
@ -74,6 +107,7 @@ inline std::tuple<const int*, int3, int3> GetRawMaskBufferAddrSizesAndStrides(
strides = {attn->max_sequence_length * attn->max_sequence_length, attn->max_sequence_length, 1};
break;
default:
LOGS_DEFAULT(FATAL) << "unsupported mask type: " << attn->mask_type;
throw std::runtime_error("unsupported mask type");
}
return {offseted_buffer, sizes, strides};
@ -82,8 +116,14 @@ inline std::tuple<const int*, int3, int3> GetRawMaskBufferAddrSizesAndStrides(
template <typename T>
struct GemmSoftmaxGemmPermuteParams : onnxruntime::rocm::tunable::OpParams {
std::string Signature() const override {
auto [m, n, k, o, batch] = GetGemmsMNKOBatch();
return MakeString("M", m, "_N", n, "_K", k, "_O", o, "_B", batch);
return MakeString(
"B", attention->batch_size,
"_S", attention->sequence_length,
"_T", attention->total_sequence_length,
"_N", attention->num_heads,
"_H", attention->head_size,
bias_buffer != nullptr ? "_B" : "_NB",
"_M", mask_index_dims.size());
}
std::tuple<int, int, int, int, int> GetGemmsMNKOBatch() const {
@ -111,10 +151,10 @@ struct GemmSoftmaxGemmPermuteParams : onnxruntime::rocm::tunable::OpParams {
// optional, mask value
const int* mask_index_buffer{nullptr};
gsl::span<const int64_t> mask_index_dims{};
TensorShapeVector mask_index_dims{};
// optional, internal
T* workspace_buffer{nullptr};
void* workspace_buffer{nullptr};
};
template <typename T>
@ -130,7 +170,7 @@ struct GemmSoftmaxGemmPermuteGenericPipeline {
params->attention->num_heads,
params->attention->sequence_length,
params->attention->total_sequence_length);
auto gemm1_out = params->workspace_buffer;
auto gemm1_out = reinterpret_cast<T*>(params->workspace_buffer);
auto softmax_out = gemm1_out + (bytes / sizeof(T));
auto gemm2_out = softmax_out + (bytes / sizeof(T));
return {gemm1_out, softmax_out, gemm2_out};
@ -143,7 +183,7 @@ struct GemmSoftmaxGemmPermuteGenericPipeline {
attn->num_heads,
attn->head_size,
attn->sequence_length,
attn->past_sequence_length);
attn->total_sequence_length);
}
inline static Status Gemm1(const GemmSoftmaxGemmPermuteParams<T>* params) {
@ -244,6 +284,247 @@ struct GemmSoftmaxGemmPermuteGenericPipeline {
}
};
namespace {
template <typename T>
struct DataTypeAdaptor {
using type = T;
};
template <>
struct DataTypeAdaptor<half> {
using type = ck::half_t;
};
template <>
struct DataTypeAdaptor<BFloat16> {
using type = ck::bhalf16_t;
};
} // namespace
template <typename T>
class GemmSoftmaxGemmPermuteTunableOp : public tunable::TunableOp<GemmSoftmaxGemmPermuteParams<T>> {
public:
GemmSoftmaxGemmPermuteTunableOp();
inline static bool IsSupportedMaskType(const AttentionParameters* attn) {
switch (attn->mask_type) {
case MASK_NONE:
case MASK_2D_DUMMY:
case MASK_2D_KEY_PADDING:
case MASK_3D_ATTENTION:
case MASK_4D_MEGATRON:
return true;
default:
return false;
}
}
inline static size_t GetWorkspaceNumBytes(const AttentionParameters* attn) {
if (!IsSupportedMaskType(attn)) {
return 0;
}
auto [buffer, sizes, strides] = GetRawMaskBufferAddrSizesAndStrides(nullptr, attn);
return sizeof(T) * sizes.x * sizes.y * sizes.z;
}
template <int VecSize, typename Converter>
__global__ static void ConvertToFilledMaskValue(
T* __restrict__ out,
const int3 out_strides,
const int* __restrict__ mask_buffer,
const int3 mask_lengths, // [B,S,T]
const int3 mask_strides,
Converter cvt) {
const int64_t global_idx = blockDim.x * blockIdx.x + threadIdx.x;
if (global_idx >= mask_lengths.x * mask_lengths.y * CeilDiv(mask_lengths.z, VecSize)) {
return;
}
const int tidx = (global_idx % CeilDiv(mask_lengths.z, VecSize)) * VecSize;
const int bs_idx = global_idx / CeilDiv(mask_lengths.z, VecSize);
const int sidx = bs_idx % mask_lengths.y;
const int bidx = bs_idx / mask_lengths.y;
int64_t in_offset = mask_strides.x * bidx + mask_strides.y * sidx + mask_strides.z * tidx;
int64_t out_offset = out_strides.x * bidx + out_strides.y * sidx + out_strides.z * tidx;
if (tidx + VecSize <= mask_lengths.z) {
using LoadT = const aligned_vector<int, VecSize>;
using StoreT = aligned_vector<T, VecSize>;
LoadT load = *reinterpret_cast<LoadT*>(mask_buffer + in_offset);
StoreT store;
#pragma unroll
for (int i = 0; i < VecSize; i++) {
store.val[i] = cvt(load.val[i]);
}
*reinterpret_cast<StoreT*>(out + out_offset) = store;
} else {
#pragma unroll
for (int i = tidx; i < mask_lengths.z; i++) {
out[out_offset + i] = cvt(mask_buffer[in_offset + i]);
}
}
}
static Status LaunchConvertToFilledMaskValue(const GemmSoftmaxGemmPermuteParams<T>* params) {
constexpr const int kThreadPerBlock = 256;
constexpr const int kVecSize = 4;
auto attn = params->attention;
auto [buffer, lengths, strides] = GetRawMaskBufferAddrSizesAndStrides(params->mask_index_buffer, attn);
int64_t total_threads = lengths.x * lengths.y * CeilDiv(lengths.z, kVecSize);
auto num_blocks = CeilDiv(total_threads, kThreadPerBlock);
auto mask_filter_value = attn->mask_filter_value;
auto cvt = [=] __device__(int v) -> T {
return v == 1 ? 0 : mask_filter_value;
};
ConvertToFilledMaskValue<kVecSize><<<num_blocks, kThreadPerBlock, 0, params->Stream()>>>(
reinterpret_cast<T*>(params->workspace_buffer), {lengths.y * lengths.z, lengths.z, 1}, // out desc
buffer, lengths, strides, // mask desc
cvt);
return HIP_CALL(hipGetLastError());
}
};
template <typename T, bool USE_BIAS, bool USE_MASK>
auto GetCKGemmSoftmaxGemmPermuteTypeStringAndOps() {
constexpr const int kNumBiasBuffer = static_cast<int>(USE_BIAS) + static_cast<int>(USE_MASK);
using Nop = ck::tensor_operation::element_wise::PassThrough;
using Acc0ElementOp = internal::PreSoftmaxAttentionScoreOp;
using CKDataType = typename DataTypeAdaptor<T>::type;
using D0DataType = typename ck::detail::tuple_concat<
std::conditional_t<USE_BIAS, ck::Tuple<CKDataType>, ck::Tuple<>>,
std::conditional_t<USE_MASK, ck::Tuple<CKDataType>, ck::Tuple<>>>::type;
constexpr static auto MaskingSpec =
ck::tensor_operation::device::MaskingSpecialization::MaskDisabled;
std::vector<std::pair<std::string, Op<GemmSoftmaxGemmPermuteParams<T>>>> ret;
for (auto&& impl : internal::GetDeviceBatchedGemmSoftmaxGemmPermuteInstances<
CKDataType, D0DataType, internal::F32, internal::PreSoftmaxAttentionScoreOp, MaskingSpec>()) {
auto type_string = impl->GetTypeString();
auto invoker = impl->MakeInvokerPointer();
auto op = [impl = std::move(impl), invoker = std::move(invoker)](
const GemmSoftmaxGemmPermuteParams<T>* params) -> Status {
if constexpr (USE_BIAS) {
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->bias_buffer == nullptr);
} else {
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->bias_buffer != nullptr);
}
if constexpr (USE_MASK) {
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(
!GemmSoftmaxGemmPermuteTunableOp<T>::IsSupportedMaskType(params->attention));
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->mask_index_buffer == nullptr);
} else {
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(params->mask_index_buffer != nullptr);
}
auto attn = params->attention;
int G0 = attn->batch_size;
int G1 = attn->num_heads;
int M = attn->sequence_length;
int N = attn->total_sequence_length;
int K = attn->head_size;
int O = attn->v_head_size;
{
auto [m, n, k, o, batch] = params->GetGemmsMNKOBatch();
ORT_ENFORCE(M == m && N == n && K == k && O == o && G0 * G1 == batch, "semantic mismatch");
ORT_ENFORCE(K == O, "inner product dimension mismatch");
}
std::vector<ck::index_t> q_buffer_lengths = {G0, G1, M, K};
std::vector<ck::index_t> q_buffer_strides = {G1 * M * K, M * K, K, 1};
std::vector<ck::index_t> k_buffer_lengths = {G0, G1, N, K};
std::vector<ck::index_t> k_buffer_strides = {G1 * N * K, N * K, K, 1}; // matrices are transposed
std::vector<ck::index_t> v_buffer_lengths = {G0, G1, O, N};
std::vector<ck::index_t> v_buffer_strides = {G1 * N * O, N * O, 1, O};
std::vector<ck::index_t> out_buffer_lengths = {G0, G1, M, O};
std::vector<ck::index_t> out_buffer_strides = {M * G1 * O, O, G1 * O, 1}; // permute 0213
std::array<void*, kNumBiasBuffer> bias_buffers{};
std::array<std::vector<ck::index_t>, kNumBiasBuffer> bias_lengths{};
std::array<std::vector<ck::index_t>, kNumBiasBuffer> bias_strides{};
if constexpr (USE_BIAS) {
bias_buffers[0] = const_cast<T*>(params->bias_buffer);
bias_lengths[0] = {G0, G1, M, N}; // BN(G0*G1), S(M), T(N)
bias_strides[0] = {G1 * M * N, M * N, N, 1};
}
if constexpr (USE_MASK) {
bias_buffers[kNumBiasBuffer - 1] = params->workspace_buffer;
bias_lengths[kNumBiasBuffer - 1] = {G0, G1, M, N}; // BN(G0*G1), S(M), T(N)
if (params->mask_index_dims.size() == 2) { // [B,T]
bias_strides[kNumBiasBuffer - 1] = {N, 0, 0, 1};
} else if (params->mask_index_dims.size() == 3) { // [B,S,T]
bias_strides[kNumBiasBuffer - 1] = {M * N, 0, N, 1};
} else if (params->mask_index_dims.size() == 4) { // [B,1,max_seq_len,max_seq_len] -->convert--> [B,S,T]
bias_strides[kNumBiasBuffer - 1] = {M * N, 0, N, 1};
} else {
ORT_ENFORCE(false, "Unreachable");
}
}
auto arg = impl->MakeArgumentPointer(
params->q_buffer, params->k_buffer, params->v_buffer, params->out_buffer,
bias_buffers, // Gemm1 bias, as attention mask
{}, // Gemm2 bias
q_buffer_lengths, q_buffer_strides,
k_buffer_lengths, k_buffer_strides,
v_buffer_lengths, v_buffer_strides,
out_buffer_lengths, out_buffer_strides,
bias_lengths, bias_strides,
{},
{},
Nop{},
Nop{},
Acc0ElementOp{params->scale},
Nop{},
Nop{});
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(!impl->IsSupportedArgument(arg.get()),
impl->GetTypeString(), " does not support ", params->Signature());
if constexpr (USE_MASK) {
ORT_RETURN_IF_ERROR(GemmSoftmaxGemmPermuteTunableOp<T>::LaunchConvertToFilledMaskValue(params));
}
invoker->Run(arg.get(), StreamConfig{params->Stream()});
return Status::OK();
};
ret.emplace_back(std::make_pair(std::move(type_string), std::move(op)));
}
return ret;
}
template <typename T>
GemmSoftmaxGemmPermuteTunableOp<T>::GemmSoftmaxGemmPermuteTunableOp() {
this->RegisterOp([](const GemmSoftmaxGemmPermuteParams<T>* params) {
return GemmSoftmaxGemmPermuteGenericPipeline<T>::Run(params, false);
});
for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps<T, /*USE_BIAS=*/false, /*USE_MASK=*/false>()) {
this->RegisterOp(std::move(op));
}
for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps<T, /*USE_BIAS=*/true, /*USE_MASK=*/false>()) {
this->RegisterOp(std::move(op));
}
for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps<T, /*USE_BIAS=*/false, /*USE_MASK=*/true>()) {
this->RegisterOp(std::move(op));
}
for (auto&& [_, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps<T, /*USE_BIAS=*/true, /*USE_MASK=*/true>()) {
this->RegisterOp(std::move(op));
}
}
} // namespace rocm
} // namespace contrib
} // namespace onnxruntime

6
onnxruntime/core/providers/rocm/math/softmax_ck.cuh
View file

@ -56,8 +56,8 @@ auto GetCKSoftmaxTypeStringAndOps() {
auto invoker = impl->MakeInvokerPointer();
auto ck_softmax_op = [impl = std::move(impl), invoker = std::move(invoker)](const SoftmaxParams<InputT, OutputT>* params) -> Status {
AccDataType alpha{1.0f};
AccDataType beta{0.0f};
double alpha{1.0f};
double beta{0.0f};
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(
params->is_log_softmax,
@ -71,7 +71,7 @@ auto GetCKSoftmaxTypeStringAndOps() {
std::vector<ck::index_t> reduce_dims{3};
auto nop = Nop{};
auto arg = impl->MakeArgumentPointer(in_lengths, in_strides, reduce_dims, &alpha, &beta,
auto arg = impl->MakeArgumentPointer(in_lengths, in_strides, reduce_dims, alpha, beta,
params->input, params->output, nop, nop);
TUNABLE_OP_RETURN_UNSUPPORTED_ARGUMENT_IF(!impl->IsSupportedArgument(arg.get()),
impl->GetTypeString(), " does not support ", params->Signature());

2
onnxruntime/python/tools/kernel_explorer/device_array.h
View file

@ -45,7 +45,7 @@ class DeviceArray {
CALL_THROW(MEMCPY(device_.get(), host_, size_ * itemsize_, MEMCPY_HOST_TO_DEVICE));
}
DeviceArray(const DeviceArray&) = default;
DeviceArray& operator=(DeviceArray&) = default;
DeviceArray& operator=(const DeviceArray&) = default;
void UpdateHostNumpyArray() {
CALL_THROW(MEMCPY(host_, device_.get(), size_ * itemsize_, MEMCPY_DEVICE_TO_HOST));

2
onnxruntime/python/tools/kernel_explorer/kernel_explorer.cc
View file

@ -8,6 +8,7 @@
#include "python/tools/kernel_explorer/kernels/rocm/fast_gelu.h"
#include "python/tools/kernel_explorer/kernels/rocm/gemm.h"
#include "python/tools/kernel_explorer/kernels/rocm/gemm_fast_gelu.h"
#include "python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.h"
#include "python/tools/kernel_explorer/kernels/rocm/skip_layer_norm.h"
#include "python/tools/kernel_explorer/kernels/rocm/softmax.h"
@ -26,6 +27,7 @@ PYBIND11_MODULE(_kernel_explorer, m) {
InitSkipLayerNorm(m);
InitGemmFastGelu(m);
InitSoftmax(m);
InitGemmSoftmaxGemmPermute(m);
#endif
m.def("is_composable_kernel_available", []() {

14
onnxruntime/python/tools/kernel_explorer/kernel_explorer_interface.h
View file

@ -58,18 +58,22 @@ class IKernelExplorer {
virtual ~IKernelExplorer() = default;
protected:
TuningContextT* TuningContext() {
if (ep_ == nullptr) {
ExecutionProvider* GetEp() {
std::call_once(ep_create_once_, [this]() {
ExecutionProviderInfo info{};
ep_ = std::make_unique<ExecutionProvider>(info);
}
this->ep_ = std::make_unique<ExecutionProvider>(info);
});
return ep_.get();
}
return static_cast<TuningContextT*>(ep_->GetTuningContext());
TuningContextT* TuningContext() {
return static_cast<TuningContextT*>(GetEp()->GetTuningContext());
}
StreamT Stream() { return stream_; }
private:
std::once_flag ep_create_once_;
std::unique_ptr<ExecutionProvider> ep_{};
StreamT stream_{0};
int repeats_{100};

375
onnxruntime/python/tools/kernel_explorer/kernels/gemm_softmax_gemm_permute_test.py Normal file
View file

@ -0,0 +1,375 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
import os
import sys
from dataclasses import dataclass
from itertools import product
import kernel_explorer as ke
import numpy as np
import pytest
from utils import dtype_to_suffix, softmax
def multinormal_distribution(num_distribution, num_element_per_dist):
arrays = []
for _ in range(num_distribution):
mean = np.random.rand() - 0.5
std = np.random.rand() + 0.5
arrays.append(np.random.normal(mean, std, (num_element_per_dist,)))
return np.array(arrays)
def get_ck_binding_name(dtype, biased: bool, masked: bool):
dtype_suffix = "_" + dtype_to_suffix(dtype)
ck_suffix = ""
if biased:
ck_suffix += "Biased"
if masked:
ck_suffix += "Masked"
ck_suffix += dtype_suffix
return "GemmSoftmaxGemmPermuteCK" + ck_suffix
dtypes = ["float16"]
batches = [1, 64]
seqlens = [128, 512]
total_seqlens = [128, 512]
num_heads = [8, 12]
head_sizes = [64]
biaseds = [False, True]
mask_dims = [0, 2, 3, 4]
def get_biased_id(biased):
return "biased" if biased else "nobias"
def get_mask_dim_id(dim):
if dim == 0:
return "nomask"
return f"mask_{dim}d"
def _test_gemm_softmax_gemm_permute(
f, dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale
):
v_head_size = head_size
q_shape = [batch, num_heads, seqlen, head_size]
k_shape = [batch, num_heads, total_seqlen, head_size]
v_shape = [batch, num_heads, total_seqlen, v_head_size]
out_shape = [batch, seqlen, num_heads, head_size]
attn_bias = None
bias_shape = [batch, num_heads, seqlen, total_seqlen] if biased else None
attn_mask = None
mask_shape = None
mask_shape_broadcasted = None
max_seqlen = None
if mask_dim != 0:
if mask_dim == 2:
mask_shape = [batch, total_seqlen]
mask_shape_broadcasted = [batch, 1, 1, total_seqlen]
elif mask_dim == 3:
mask_shape = [batch, seqlen, total_seqlen]
mask_shape_broadcasted = [batch, 1, seqlen, total_seqlen]
elif mask_dim == 4:
max_seqlen = ((seqlen - 1) // 1024 + 1) * 1024 # round up to multiple of 1024
mask_shape = [batch, 1, max_seqlen, max_seqlen]
else:
raise ValueError
np.random.seed(42)
Q = multinormal_distribution(np.prod(q_shape[:-1]), q_shape[-1]).reshape(q_shape).astype(np.float64)
K = multinormal_distribution(np.prod(k_shape[:-1]), k_shape[-1]).reshape(k_shape).astype(np.float64)
V = multinormal_distribution(np.prod(v_shape[:-1]), v_shape[-1]).reshape(v_shape).astype(np.float64)
if bias_shape is not None:
attn_bias = np.random.uniform(-0.5, 0.5, size=bias_shape)
if mask_shape is not None:
attn_mask = (np.random.randint(0, 100, size=mask_shape) < 95).astype(np.int32)
pre_softmax_attn_scores = Q @ np.swapaxes(K, 2, 3)
pre_softmax_attn_scores = pre_softmax_attn_scores * scale
if attn_bias is not None:
pre_softmax_attn_scores = pre_softmax_attn_scores + attn_bias
if attn_mask is not None:
filter_value = -10000.0
if mask_dim == 4:
# equivalent to past_sequence_length = max_sequence_length - seqlen
converted_mask = (1 - attn_mask[:, :, -seqlen:, :total_seqlen]) * filter_value
else:
converted_mask = (1 - attn_mask.reshape(mask_shape_broadcasted)) * filter_value
pre_softmax_attn_scores = pre_softmax_attn_scores + converted_mask
attn_scores = softmax(pre_softmax_attn_scores, axis=-1)
attn = attn_scores @ V
ref = np.swapaxes(attn, 2, 1) # permute 0213
out = np.empty(out_shape, dtype=dtype)
host_Q = Q.astype(dtype)
host_K = K.astype(dtype)
host_V = V.astype(dtype)
host_attn_bias = attn_bias.astype(dtype) if attn_bias is not None else None
dev_Q = ke.DeviceArray(host_Q)
dev_K = ke.DeviceArray(host_K)
dev_V = ke.DeviceArray(host_V)
dev_out = ke.DeviceArray(out)
dev_attn_bias = ke.DeviceArray(host_attn_bias) if host_attn_bias is not None else None
dev_attn_mask = ke.DeviceArray(attn_mask) if attn_mask is not None else None
my_gemm_softmax_gemm_permute = f(
batch,
seqlen,
total_seqlen,
max_seqlen,
num_heads,
head_size,
mask_dim,
scale,
dev_Q,
dev_K,
dev_V,
dev_attn_bias,
dev_attn_mask,
dev_out,
)
print() # write an empty line in case pytest ... -s -v
failures = {}
for impl in my_gemm_softmax_gemm_permute.ListOps():
if not my_gemm_softmax_gemm_permute.SelectOp(impl):
print("Unsupport", impl)
continue
print(" Support", impl)
my_gemm_softmax_gemm_permute.Run()
dev_out.UpdateHostNumpyArray()
try:
is_strict = int(os.environ.get("KERNEL_EXPLORER_STRICT_TEST", "0"))
if is_strict:
# NOTE: this will always fail, just for manual checking with:
# KERNEL_EXPLORER_STRICT_TEST=1 pytest ... -s -v
np.testing.assert_allclose(out, ref)
else:
is_zero_tol, atol, rtol = 1e-3, 1e-2, 1e-2
not_close_to_zeros = np.abs(ref) > is_zero_tol
np.testing.assert_allclose(out[not_close_to_zeros], ref[not_close_to_zeros], atol=atol, rtol=rtol)
except Exception as err:
header = "*" * 30 + impl + "*" * 30
print(header)
print(err)
print("*" * len(header))
failures[impl] = str(err)
if failures:
raise Exception(failures)
@pytest.mark.parametrize("mask_dim", mask_dims, ids=get_mask_dim_id)
@pytest.mark.parametrize("biased", biaseds, ids=get_biased_id)
@pytest.mark.parametrize("head_size", head_sizes)
@pytest.mark.parametrize("nhead", num_heads)
@pytest.mark.parametrize("total_seqlen", total_seqlens)
@pytest.mark.parametrize("seqlen", seqlens)
@pytest.mark.parametrize("batch", [16])
@pytest.mark.parametrize("dtype", dtypes)
def test_gemm_softmax_gemm_permute_generic(dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim):
f = getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype))
scale = 1.0 / np.sqrt(head_size)
_test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale)
@pytest.mark.parametrize("mask_dim", mask_dims, ids=get_mask_dim_id)
@pytest.mark.parametrize("biased", biaseds, ids=get_biased_id)
@pytest.mark.parametrize("head_size", head_sizes)
@pytest.mark.parametrize("nhead", num_heads)
@pytest.mark.parametrize("total_seqlen", total_seqlens)
@pytest.mark.parametrize("seqlen", seqlens)
@pytest.mark.parametrize("batch", batches)
@pytest.mark.parametrize("dtype", dtypes)
def test_gemm_softmax_gemm_permute_ck(dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim):
f = getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0))
scale = 1.0 / np.sqrt(head_size)
_test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale)
@pytest.mark.parametrize("mask_dim", [2], ids=get_mask_dim_id)
@pytest.mark.parametrize("biased", [False], ids=get_biased_id)
@pytest.mark.parametrize("head_size", [64])
@pytest.mark.parametrize("nhead", [8])
@pytest.mark.parametrize("total_seqlen", [128])
@pytest.mark.parametrize("seqlen", [64])
@pytest.mark.parametrize("batch", [16])
@pytest.mark.parametrize("dtype", ["float16"])
def test_gemm_softmax_gemm_permute_tunable(dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim):
f = getattr(ke, "GemmSoftmaxGemmPermuteTunable_" + dtype_to_suffix(dtype))
scale = 1.0 / np.sqrt(head_size)
_test_gemm_softmax_gemm_permute(f, dtype, batch, seqlen, total_seqlen, nhead, head_size, biased, mask_dim, scale)
@dataclass
class GemmSoftmaxGemmPermuteMetric(ke.ComputeMetric):
batch: int
seqlen: int
total_seqlen: int
num_heads: int
head_size: int
biased: bool
mask_dim: int
def report(self):
bias_str = " biased" if self.biased else ""
mask_str = f" mask_{self.mask_dim}d" if self.mask_dim != 0 else ""
common = (
f"{self.dtype} B={self.batch} S={self.seqlen} T={self.total_seqlen} "
+ f"N={self.num_heads} H={self.head_size}{bias_str}{mask_str}, "
+ f"{self.name}"
)
if self.duration <= 0:
return "not supported " + common
return f"{self.duration:>6.2f} us {self.tflops:>5.2f} tflops " + common
def profile_gemm_softmax_gemm_permute_func(
f, dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale
):
v_head_size = head_size
q_shape = [batch, num_heads, seqlen, head_size]
k_shape = [batch, num_heads, total_seqlen, head_size]
v_shape = [batch, num_heads, total_seqlen, v_head_size]
out_shape = [batch, seqlen, num_heads, head_size]
attn_bias = None
bias_shape = [batch, num_heads, seqlen, total_seqlen] if biased else None
attn_mask = None
mask_shape = None
max_seqlen = None
if mask_dim != 0:
if mask_dim == 2:
mask_shape = [batch, total_seqlen]
elif mask_dim == 3:
mask_shape = [batch, seqlen, total_seqlen]
elif mask_dim == 4:
max_seqlen = ((seqlen - 1) // 1024 + 1) * 1024 # round up to multiple of 1024
mask_shape = [batch, 1, max_seqlen, max_seqlen]
else:
raise ValueError
np.random.seed(42)
Q = multinormal_distribution(np.prod(q_shape[:-1]), q_shape[-1]).reshape(q_shape).astype(np.float64)
K = multinormal_distribution(np.prod(k_shape[:-1]), k_shape[-1]).reshape(k_shape).astype(np.float64)
V = multinormal_distribution(np.prod(v_shape[:-1]), v_shape[-1]).reshape(v_shape).astype(np.float64)
if bias_shape is not None:
attn_bias = np.random.uniform(-2, 2, size=bias_shape)
if mask_shape is not None:
attn_mask = (np.random.randint(0, 100, size=mask_shape) < 95).astype(np.int32)
out = np.empty(out_shape, dtype=dtype)
host_Q = Q.astype(dtype)
host_K = K.astype(dtype)
host_V = V.astype(dtype)
host_attn_bias = attn_bias.astype(dtype) if attn_bias is not None else None
dev_Q = ke.DeviceArray(host_Q)
dev_K = ke.DeviceArray(host_K)
dev_V = ke.DeviceArray(host_V)
dev_out = ke.DeviceArray(out)
dev_attn_bias = ke.DeviceArray(host_attn_bias) if host_attn_bias is not None else None
dev_attn_mask = ke.DeviceArray(attn_mask) if attn_mask is not None else None
my_gemm_softmax_gemm_permute = f(
batch,
seqlen,
total_seqlen,
max_seqlen,
num_heads,
head_size,
mask_dim,
scale,
dev_Q,
dev_K,
dev_V,
dev_attn_bias,
dev_attn_mask,
dev_out,
)
for impl in my_gemm_softmax_gemm_permute.ListOps():
duration_ms = -1
if my_gemm_softmax_gemm_permute.SelectOp(impl):
duration_ms = my_gemm_softmax_gemm_permute.Profile()
m, n, k, o, gemm_batch = seqlen, total_seqlen, head_size, head_size, batch * num_heads
flops_per_batch = m * n * k * 2 + m * n * o * 2
flops_count_bias_and_softmax = True # set to false to be aligned with ck
if flops_count_bias_and_softmax:
flops_per_batch += 2 * n + 1
if flops_count_bias_and_softmax and attn_bias is not None:
flops_per_batch += m * n
if flops_count_bias_and_softmax and attn_mask is not None:
flops_per_batch += m * n
flops = flops_per_batch * gemm_batch
ke.report(
GemmSoftmaxGemmPermuteMetric(
impl, dtype, duration_ms, flops, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim
)
)
def profile_with_args(dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale, *, sort=False):
with ke.benchmark(sort):
args = (dtype, batch, seqlen, total_seqlen, num_heads, head_size, biased, mask_dim, scale)
profile_gemm_softmax_gemm_permute_func(
getattr(ke, "GemmSoftmaxGemmPermuteGeneric_" + dtype_to_suffix(dtype)), *args
)
profile_gemm_softmax_gemm_permute_func(getattr(ke, get_ck_binding_name(dtype, biased, mask_dim != 0)), *args)
profile_gemm_softmax_gemm_permute_func(
getattr(ke, "GemmSoftmaxGemmPermuteTunable_" + dtype_to_suffix(dtype)), *args
)
def profile():
for args in product(dtypes, batches, seqlens, total_seqlens, num_heads, head_sizes, biaseds, mask_dims):
profile_with_args(*args, scale=0.125, sort=True)
print()
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
group = parser.add_argument_group("profile with args")
group.add_argument("--sort", action="store_true")
group.add_argument("dtype", choices=dtypes)
group.add_argument("batch", type=int)
group.add_argument("seqlen", type=int)
group.add_argument("total_seqlen", type=int)
group.add_argument("num_heads", type=int)
group.add_argument("head_size", type=int)
group.add_argument("biased", type=int, choices=[0, 1], default=0)
group.add_argument("mask_dim", type=int, choices=[0, 2, 3, 4], default=2, help="0 for mask disabled")
group.add_argument("scale", type=float, default=None, help="default to 1.0/sqrt(head_size)")
if len(sys.argv) == 1:
profile()
else:
args = parser.parse_args()
profile_with_args(
args.dtype,
args.batch,
args.seqlen,
args.total_seqlen,
args.num_heads,
args.head_size,
args.biased,
args.mask_dim,
args.scale,
sort=args.sort,
)

298
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.cu Normal file
View file

@ -0,0 +1,298 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.h"
#include "pybind11/stl.h"
#include "contrib_ops/rocm/bert/batched_gemm_softmax_gemm_permute_pipelines.cuh"
#include "core/providers/rocm/tunable/rocm_tunable.h"
#include "python/tools/kernel_explorer/device_array.h"
#include "python/tools/kernel_explorer/kernel_explorer_interface.h"
#include <vector>
namespace py = pybind11;
using namespace onnxruntime::contrib::rocm;
namespace onnxruntime {
template <typename T>
class IGemmSoftmaxGemmPermuteKernelExplorer : public IKernelExplorer {
public:
IGemmSoftmaxGemmPermuteKernelExplorer(
int64_t batch,
int64_t seqlen,
int64_t total_seqlen,
std::optional<int64_t> max_seqlen,
int64_t num_heads,
int64_t head_size,
int64_t mask_dim,
double scale,
DeviceArray& Q,
DeviceArray& K,
DeviceArray& V,
std::optional<DeviceArray>& attn_bias,
std::optional<DeviceArray>& attn_mask,
DeviceArray& out) {
ROCBLAS_CALL_THROW(rocblas_create_handle(&rocblas_handle_));
attn_.batch_size = batch;
attn_.sequence_length = seqlen;
attn_.kv_sequence_length = seqlen; // NOTE: not used
attn_.past_sequence_length = 0;
attn_.original_past_sequence_length = 0; // NOTE: not used
attn_.total_sequence_length = total_seqlen;
attn_.max_sequence_length = 0;
attn_.hidden_size = num_heads * head_size;
attn_.head_size = head_size;
attn_.v_hidden_size = attn_.hidden_size; // Q,K,V hidden size must agree now
attn_.v_head_size = attn_.head_size; // Q,K,V hidden size must agree now
attn_.num_heads = num_heads;
attn_.is_unidirectional = false;
attn_.past_present_share_buffer = false;
attn_.do_rotary = false;
attn_.mask_filter_value = -10000.0f;
attn_.scale = scale;
if (mask_dim == 0) {
attn_.mask_type = contrib::MASK_NONE;
} else if (mask_dim == 2) {
attn_.mask_type = contrib::MASK_2D_KEY_PADDING;
} else if (mask_dim == 3) {
attn_.mask_type = contrib::MASK_3D_ATTENTION;
} else if (mask_dim == 4) {
attn_.mask_type = contrib::MASK_4D_MEGATRON;
} else {
ORT_ENFORCE(false, "mask type not supported");
}
device_prop = GetEp()->GetDeviceProp();
params_.tuning_ctx = TuningContext();
params_.stream = Stream();
params_.handle = rocblas_handle_;
params_.attention = &attn_;
params_.device_prop = &device_prop;
params_.scale = scale;
params_.q_buffer = reinterpret_cast<T*>(Q.ptr());
params_.k_buffer = reinterpret_cast<T*>(K.ptr());
params_.v_buffer = reinterpret_cast<T*>(V.ptr());
if (attn_bias.has_value()) {
params_.bias_buffer = reinterpret_cast<T*>(attn_bias->ptr());
}
if (attn_mask.has_value()) {
params_.mask_index_buffer = reinterpret_cast<int*>(attn_mask->ptr());
if (mask_dim == 2) {
params_.mask_index_dims = {batch, total_seqlen};
} else if (mask_dim == 3) {
params_.mask_index_dims = {batch, seqlen, total_seqlen};
} else if (mask_dim == 4) {
ORT_ENFORCE(max_seqlen.has_value());
attn_.max_sequence_length = max_seqlen.value();
ORT_ENFORCE(attn_.max_sequence_length >= seqlen);
attn_.past_sequence_length = attn_.max_sequence_length - seqlen;
params_.mask_index_dims = {batch, 1, attn_.max_sequence_length, attn_.max_sequence_length};
}
}
params_.out_buffer = reinterpret_cast<T*>(out.ptr());
}
~IGemmSoftmaxGemmPermuteKernelExplorer() {
ROCBLAS_CALL_THROW(rocblas_destroy_handle(rocblas_handle_));
}
void SetWorkspace(size_t num_bytes) {
void* ptr;
HIP_CALL_THROW(hipMalloc(&ptr, num_bytes));
workspace_.reset(ptr, [](void* ptr) { HIP_CALL_THROW(hipFree(ptr)); });
params_.workspace_buffer = reinterpret_cast<T*>(workspace_.get());
}
protected:
using ParamsT = contrib::rocm::GemmSoftmaxGemmPermuteParams<T>;
rocblas_handle rocblas_handle_;
hipDeviceProp_t device_prop;
contrib::AttentionParameters attn_;
ParamsT params_;
std::shared_ptr<void> workspace_;
};
// The pipeline composed from rocblas api calls and kernel launches.
template <typename T>
class GemmSoftmaxGemmPermuteGeneric : public IGemmSoftmaxGemmPermuteKernelExplorer<T> {
public:
GemmSoftmaxGemmPermuteGeneric(
int64_t batch,
int64_t seqlen,
int64_t total_seqlen,
std::optional<int64_t> max_seqlen,
int64_t num_heads,
int64_t head_size,
int64_t mask_dim,
double scale,
DeviceArray& Q,
DeviceArray& K,
DeviceArray& V,
std::optional<DeviceArray>& attn_bias,
std::optional<DeviceArray>& attn_mask,
DeviceArray& out)
: IGemmSoftmaxGemmPermuteKernelExplorer<T>(batch, seqlen, total_seqlen, max_seqlen,
num_heads, head_size, mask_dim, scale,
Q, K, V, attn_bias, attn_mask, out) {
this->SetWorkspace(GemmSoftmaxGemmPermuteGenericPipeline<T>::GetWorkspaceNumBytes(&this->attn_));
}
std::vector<std::string> ListOps() const {
return {"Generic"};
}
bool SelectOp(const std::string&) {
return true;
}
void Run() override {
ORT_THROW_IF_ERROR(GemmSoftmaxGemmPermuteGenericPipeline<T>::Run(
&this->params_, /*use_persistent_softmax=*/false));
}
};
template <typename T, bool USE_BIAS, bool USE_MASK>
class GemmSoftmaxGemmPermuteCK : public IGemmSoftmaxGemmPermuteKernelExplorer<T> {
public:
GemmSoftmaxGemmPermuteCK(
int64_t batch,
int64_t seqlen,
int64_t total_seqlen,
std::optional<int64_t> max_seqlen,
int64_t num_heads,
int64_t head_size,
int64_t mask_dim,
double scale,
DeviceArray& Q,
DeviceArray& K,
DeviceArray& V,
std::optional<DeviceArray>& attn_bias,
std::optional<DeviceArray>& attn_mask,
DeviceArray& out)
: IGemmSoftmaxGemmPermuteKernelExplorer<T>(batch, seqlen, total_seqlen, max_seqlen,
num_heads, head_size, mask_dim, scale,
Q, K, V, attn_bias, attn_mask, out) {
this->SetWorkspace(GemmSoftmaxGemmPermuteTunableOp<T>::GetWorkspaceNumBytes(&this->attn_));
for (auto&& [ts, op] : GetCKGemmSoftmaxGemmPermuteTypeStringAndOps<T, USE_BIAS, USE_MASK>()) {
type_strings_.emplace_back(std::move(ts));
ops_.emplace_back(std::move(op));
}
}
std::vector<std::string> ListOps() const {
return type_strings_;
}
bool SelectOp(const std::string& name) {
for (size_t i = 0; i < ops_.size(); i++) {
if (type_strings_[i] == name) {
selected_op_ = i;
Status status = ops_[i].IsSupported(&this->params_);
return status.IsOK();
}
}
ORT_THROW("Cannot find implementation ", name);
}
void Run() override {
ORT_THROW_IF_ERROR(ops_[selected_op_](&this->params_));
}
private:
using ParamsT = typename IGemmSoftmaxGemmPermuteKernelExplorer<T>::ParamsT;
using OpT = Op<ParamsT>;
std::vector<OpT> ops_;
std::vector<std::string> type_strings_;
size_t selected_op_{};
};
// The pipeline composed from rocblas api calls and kernel launches.
template <typename T>
class GemmSoftmaxGemmPermuteTunable : public IGemmSoftmaxGemmPermuteKernelExplorer<T> {
public:
GemmSoftmaxGemmPermuteTunable(
int64_t batch,
int64_t seqlen,
int64_t total_seqlen,
std::optional<int64_t> max_seqlen,
int64_t num_heads,
int64_t head_size,
int64_t mask_dim,
double scale,
DeviceArray& Q,
DeviceArray& K,
DeviceArray& V,
std::optional<DeviceArray>& attn_bias,
std::optional<DeviceArray>& attn_mask,
DeviceArray& out)
: IGemmSoftmaxGemmPermuteKernelExplorer<T>(batch, seqlen, total_seqlen, max_seqlen,
num_heads, head_size, mask_dim, scale,
Q, K, V, attn_bias, attn_mask, out) {
this->SetWorkspace(std::max(
GemmSoftmaxGemmPermuteGenericPipeline<T>::GetWorkspaceNumBytes(&this->attn_),
GemmSoftmaxGemmPermuteTunableOp<T>::GetWorkspaceNumBytes(&this->attn_)));
this->params_.TuningContext()->EnableTunableOp();
}
std::vector<std::string> ListOps() const {
return {"Tunable"};
}
bool SelectOp(const std::string&) {
return true;
}
void Run() override {
ORT_THROW_IF_ERROR(GemmSoftmaxGemmPermuteTunableOp<T>{}(&this->params_));
}
};
#define REGISTER_COMMON(name, type, ...) \
py::class_<type<__VA_ARGS__>>(m, name) \
.def(py::init<int64_t, int64_t, int64_t, std::optional<int64_t>, int64_t, int64_t, int64_t, \
float, \
DeviceArray&, \
DeviceArray&, \
DeviceArray&, \
std::optional<DeviceArray>&, \
std::optional<DeviceArray>&, \
DeviceArray&>()) \
.def("SetRepeats", &type<__VA_ARGS__>::SetRepeats) \
.def("Run", &type<__VA_ARGS__>::Run) \
.def("Profile", &type<__VA_ARGS__>::Profile) \
.def("ListOps", &type<__VA_ARGS__>::ListOps) \
.def("SelectOp", &type<__VA_ARGS__>::SelectOp);
#define REGISTER_GENERIC(dtype) \
REGISTER_COMMON("GemmSoftmaxGemmPermuteGeneric_" #dtype, GemmSoftmaxGemmPermuteGeneric, dtype)
#define REGISTER_CK(dtype, biased, masked, mask_bias_suffix) \
REGISTER_COMMON( \
"GemmSoftmaxGemmPermuteCK" mask_bias_suffix "_" #dtype, GemmSoftmaxGemmPermuteCK, dtype, biased, masked)
#define REGISTER_TUNABLE(dtype) \
REGISTER_COMMON("GemmSoftmaxGemmPermuteTunable_" #dtype, GemmSoftmaxGemmPermuteTunable, dtype)
void InitGemmSoftmaxGemmPermute(py::module m) {
REGISTER_GENERIC(half);
REGISTER_CK(half, false, false, "");
REGISTER_CK(half, true, false, "Biased");
REGISTER_CK(half, false, true, "Masked");
REGISTER_CK(half, true, true, "BiasedMasked");
REGISTER_TUNABLE(half);
}
} // namespace onnxruntime

14
onnxruntime/python/tools/kernel_explorer/kernels/rocm/gemm_softmax_gemm_permute.h Normal file
View file

@ -0,0 +1,14 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <pybind11/pybind11.h>
namespace py = pybind11;
namespace onnxruntime {
void InitGemmSoftmaxGemmPermute(py::module mod);
}

11
onnxruntime/python/tools/kernel_explorer/kernels/softmax_test.py
View file

@ -11,7 +11,7 @@ from itertools import product
import kernel_explorer as ke
import numpy as np
import pytest
from utils import dtype_to_bytes, dtype_to_suffix
from utils import dtype_to_bytes, dtype_to_suffix, softmax
def get_test_sizes():
@ -29,13 +29,6 @@ def dtype_to_funcs(dtype):
return type_map[dtype]
def softmax(x, is_log_softmax):
x = x - np.max(x, axis=-1, keepdims=1)
if is_log_softmax:
return x - np.log(np.sum(np.exp(x), axis=-1, keepdims=1))
return (np.exp(x)) / np.sum(np.exp(x), axis=-1, keepdims=1)
def _test_softmax(batch_count, softmax_elements, is_log_softmax, dtype, func):
np.random.seed(0)
x = np.random.rand(batch_count, softmax_elements).astype(dtype)
@ -43,7 +36,7 @@ def _test_softmax(batch_count, softmax_elements, is_log_softmax, dtype, func):
x_d = ke.DeviceArray(x)
y_d = ke.DeviceArray(y)
y_ref = softmax(x, is_log_softmax)
y_ref = softmax(x, is_log_softmax=is_log_softmax)
softmax_func = getattr(ke, func)
softmax_op = softmax_func(

7
onnxruntime/python/tools/kernel_explorer/kernels/utils.py
View file

@ -93,3 +93,10 @@ def get_gemm_basic_sizes(full=True):
# ck has various impls to be tested, use the full basic cases will result too many cases to test.
# So we use a reduced combination here.
return list(product([1, 4, 127, 133], [3, 16, 128], [3, 129, 1024]))
def softmax(x, *, is_log_softmax=False, axis=-1):
x = x - np.max(x, axis=axis, keepdims=1)
if is_log_softmax:
return x - np.log(np.sum(np.exp(x), axis=axis, keepdims=1))
return (np.exp(x)) / np.sum(np.exp(x), axis=axis, keepdims=1)