Repository: microsoft/onnxruntime
File: onnxruntime/core/providers/cpu/llm/attention.h (C/C++)

Code before:
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/platform/threadpool.h"
#include "core/providers/cpu/llm/attention_helper.h"

namespace onnxruntime {

// This value is used to mask out a value from the input as ``Softmax(-infinity, ...) = 0``.
// If the mask is added, -infinity + x = -infinity.
// inifinity is replaced by lowest() because softmax implemented in MLAS
// is expected an input with only infinity values.
// This change assumes that lowest() is sufficiently low enough to not impact the result.
template <typename T>
inline T negative_infinity() {
  return std::numeric_limits<T>::lowest();
}

template <>
inline MLFloat16 negative_infinity() {
  return MLFloat16(negative_infinity<float>());
}

template <typename T>
class AttentionBase : public OpKernel {
 public:
  AttentionBase(const OpKernelInfo& info) : OpKernel(info) {}

  Status ApplyAttention(OpKernelContext* context,
                        const T* Q,                                              // Q data with shape BxNxSxH
                        const T* K,                                              // K data with shape BxNxLxH
                        const T* V,                                              // V value with size BxNxLxH_v
                        const Tensor* mask_index,                                // mask index. nullptr if no mask or its size is B
                        const Tensor* past_key,                                  // past K input tensor (if not using past state)
                        const Tensor* past_value,                                // past V input tensor (if not using past state)
                        Tensor* output,                                          // output tensor

Review comment (suggestion):
Shall we modify MLAS to handle -inf instead? I think exp(x) shall handle -inf no matter it is MLAS or not.

Code after:
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

#pragma once
#include "core/common/common.h"
#include "core/framework/op_kernel.h"
#include "core/platform/threadpool.h"
#include "core/providers/cpu/llm/attention_helper.h"

namespace onnxruntime {

// This value is used to mask out a value from the input as ``Softmax(-infinity, ...) = 0``.
// If the mask is added, -infinity + x = -infinity.
// inifinity is replaced by lowest() because softmax implemented in MLAS
// is expected an input with only infinity values.
// This change assumes that lowest() is sufficiently low enough to not impact the result.
template <typename T>
inline T mask_filter_value() {
  return std::numeric_limits<T>::lowest();
}

template <>
inline MLFloat16 mask_filter_value() {
  return MLFloat16::MinValue;
}

template <typename T>
class AttentionBase : public OpKernel {
 public:
  AttentionBase(const OpKernelInfo& info) : OpKernel(info) {}

  Status ApplyAttention(OpKernelContext* context,
                        const T* Q,                                              // Q data with shape BxNxSxH
                        const T* K,                                              // K data with shape BxNxLxH
                        const T* V,                                              // V value with size BxNxLxH_v
                        const Tensor* mask_index,                                // mask index. nullptr if no mask or its size is B
                        const Tensor* past_key,                                  // past K input tensor (if not using past state)
                        const Tensor* past_value,                                // past V input tensor (if not using past state)
                        Tensor* output,                                          // output tensor