saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make ROCm Attention biased+masked and biased+nomask scaling logic consistent (#14976)

Browse source 
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
This commit is contained in:
cloudhan 2023-03-09 23:37:50 +08:00 committed by GitHub
parent f83923d5df
commit 51b67fa15c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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) {