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