From 51b67fa15c703ea460aeef096745c3d591658532 Mon Sep 17 00:00:00 2001 From: cloudhan Date: Thu, 9 Mar 2023 23:37:50 +0800 Subject: [PATCH] Make ROCm Attention biased+masked and biased+nomask scaling logic consistent (#14976) The biased+masked and biased+nomask have different scaling logic in current ROCm implementation Currently, biased + masked: (QK'+ bias) * scale + convert(mask) biased + nomask: QK' * scale + bias which is not correct. What we want is QK' * scale [+ bias] That is, bias should not be scaled. This effectively follows https://github.com/microsoft/onnxruntime/pull/14517/files?w=1#diff-e4768ce15a73499f584f9cd7d71adcb1ff2ed8d68ad7e496723a4775cbc35e33 --- onnxruntime/contrib_ops/rocm/bert/attention_softmax.h | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h b/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h index 8e47a62403..8a34e29e10 100644 --- a/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h +++ b/onnxruntime/contrib_ops/rocm/bert/attention_softmax.h @@ -205,11 +205,7 @@ __global__ void SoftmaxWithRawMaskSmallKernel( // to avoid the performance impact from using the valid_items interface. float thread_data = -ROCMRT_INF_F; if (threadIdx.x < all_sequence_length) { - if (add_before_softmax == nullptr) { - thread_data = float(input[index]) * rsqrt_head_size; - } else { - thread_data = float(input[index] + add_before_softmax[index]) * rsqrt_head_size; - } + thread_data = float(input[index]) * rsqrt_head_size; const int sequence_index = blockIdx.x % sequence_length; if (is_unidirectional) { @@ -234,6 +230,10 @@ __global__ void SoftmaxWithRawMaskSmallKernel( thread_data = -ROCMRT_INF_F; } } + + if (add_before_softmax != nullptr) { + thread_data += float(add_before_softmax[index]); + } } if (skip_softmax) {