diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index f84d7593141f2..1ebe6d2e2f250 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -2387,27 +2387,48 @@ This version of the operator has been available since version 1 of the 'com.micr
### **com.microsoft.GatedAdd**
- Adds one tensor to another tensor scaled by a per-row gate:
+ Adds one tensor to another tensor scaled by a per-row gate. X and Y have shape (..., C). The
+ leading dimensions of gate must match the leading dimensions of X (there is no broadcasting over
+ them); the optional gate_weight input selects how the per-row gate is obtained.
+
+ Plain mode (gate_weight is not provided): gate has shape (..., 1) and is used directly as the
+ gate, broadcast over C. No sigmoid is applied.
output = X + round_to_T(Y * gate)
- X and Y have shape (..., C), and gate has shape (..., 1). The gate is broadcast
- over C. For reduced-precision types, the product is rounded to T before the add,
- matching separate ONNX Mul and Add operators.
+ For reduced-precision types, the product is rounded to T before the add, matching separate ONNX
+ Mul and Add operators. This mode supports float, float16 and bfloat16, subject to execution provider kernel support.
+
+ Projected mode (gate_weight is provided): gate holds an activation with shape (..., K) and
+ gate_weight has shape (K, 1). For every row, the gate is computed by projecting the activation
+ with gate_weight and applying a sigmoid:
+
+ p = round_to_bfloat16(sum_k(gate[..., k] * gate_weight[k]))
+ s = round_to_bfloat16(sigmoid(p))
+ output = round_to_bfloat16(X + round_to_bfloat16(Y * s))
+
+ The dot product, the sigmoid and the multiplication and addition are evaluated in FP32 (the
+ sigmoid is always computed in FP32, and the order of the FP32 accumulation over K is not
+ specified); the result is rounded to bfloat16 after the projection, after the sigmoid, after the
+ multiplication and after the addition. This mode is only implemented by the CUDA execution
+ provider for bfloat16 tensors. Other execution providers and types return NOT_IMPLEMENTED when
+ gate_weight is provided.
#### Version
This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
-#### Inputs
+#### Inputs (3 - 4)
- X : T
- Unscaled input with shape (..., C).
- Y : T
-- Input scaled by gate, with the same shape as X.
+- Input scaled by the gate, with the same shape as X.
- gate : T
-- Per-row gate with shape (..., 1).
+- In plain mode, the per-row gate with shape (..., 1). When gate_weight is provided, the activation with shape (..., K) that is projected with gate_weight and passed through a sigmoid to produce the per-row gate. Its leading dimensions must match X.
+- gate_weight (optional) : T
+- Optional projection weight with shape (K, 1), where K is the last dimension of gate. When present, the gate is sigmoid(gate . gate_weight) computed per row in FP32 (projected mode). Only supported by the CUDA execution provider with bfloat16.
#### Outputs
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 2497a7e6fd31a..fc416deb384c4 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -600,7 +600,7 @@ The **OpSet Version** column uses the following notation:
|FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedGemm|*in* A:**T**
*in* B:**T**
*in* C:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float)|
-|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
+|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*in* gate_weight:**T**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
|GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
|GatherBlockQuantized|*in* data:**T1**
*in* indices:**Tind**
*in* scales:**T2**
*in* zero_points:**T1**
*out* output:**T2**|1+|**T1** = tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)
**T2** = tensor(float), tensor(float16)
**Tind** = tensor(int32), tensor(int64)|
|GatherND|*in* data:**T**
*in* indices:**Tind**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)
**Tind** = tensor(int32), tensor(int64)|
@@ -686,10 +686,10 @@ The **OpSet Version** column uses the following notation:
|||[7, 12]|**T** = tensor(double), tensor(float), tensor(float16), tensor(int32), tensor(int64), tensor(uint32), tensor(uint64)|
|Affine|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|And|*in* A:**T**
*in* B:**T**
*out* C:**T1**|7+|**T** = tensor(bool)
**T1** = tensor(bool)|
-|ArgMax|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(float16)|
+|ArgMax|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|||12|**T** = tensor(double), tensor(float), tensor(float16)|
|||[1, 11]|**T** = tensor(double), tensor(float), tensor(float16)|
-|ArgMin|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(float16)|
+|ArgMin|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|||12|**T** = tensor(double), tensor(float), tensor(float16)|
|||[1, 11]|**T** = tensor(double), tensor(float), tensor(float16)|
|Attention|*in* Q:**T1**
*in* K:**T1**
*in* V:**T2**
*in* attn_mask:**U**
*in* past_key:**T1**
*in* past_value:**T2**
*in* nonpad_kv_seqlen:**tensor(int64)**
*out* Y:**T1**
*out* present_key:**T1**
*out* present_value:**T2**
*out* qk_matmul_output:**T1**
or
*in* Q:**T1**
*in* K:**T1**
*in* V:**T2**
*in* attn_mask:**U**
*in* past_key:**T1**
*in* past_value:**T2**
*out* Y:**T1**
*out* present_key:**T1**
*out* present_value:**T2**
*out* qk_matmul_output:**T1**|24+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(bfloat16), tensor(float), tensor(float16)
**U** = tensor(bfloat16), tensor(bool), tensor(float), tensor(float16)|
@@ -1119,7 +1119,7 @@ The **OpSet Version** column uses the following notation:
|FastGelu|*in* X:**T**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
-|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
+|GatedAdd|*in* X:**T**
*in* Y:**T**
*in* gate:**T**
*in* gate_weight:**T**
*out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|GatedDeltaNet|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* cu_seqlens:**TI**
*in* decay:**TS**
*in* beta:**TS**
*in* initial_state:**TS**
*in* a_log:**TS**
*in* dt_bias:**TS**
*in* capture_count:**TI**
*in* state_update_active:**TI**
*out* output:**T**
*out* final_state:**TS**
*out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)
**TI** = tensor(int32)
**TS** = tensor(float)|
|GatedRMSNorm|*in* X:**T**
*in* scale:**T**
*in* gate:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
|GatedRelativePositionBias|*in* query_layer:**T**
*in* query_bias:**T**
*in* rel_pos:**T**
*in* weight:**T**
*in* bias:**T**
*in* eco_a:**T**
*in* token_offset:**M**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
diff --git a/onnxruntime/contrib_ops/cpu/bert/gated_add.cc b/onnxruntime/contrib_ops/cpu/bert/gated_add.cc
index 48b643da3bc20..8c79885a3ee33 100644
--- a/onnxruntime/contrib_ops/cpu/bert/gated_add.cc
+++ b/onnxruntime/contrib_ops/cpu/bert/gated_add.cc
@@ -46,6 +46,10 @@ Status GatedAdd::Compute(OpKernelContext* context) const {
const Tensor* x = context->Input(0);
const Tensor* y = context->Input(1);
const Tensor* gate = context->Input(2);
+ const Tensor* gate_weight = context->Input(3);
+ if (gate_weight != nullptr) {
+ return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "GatedAdd gate_weight currently supports BFloat16 only");
+ }
const TensorShape& shape = x->Shape();
ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1");
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc
index 9480f08c43d30..3f5b3349d4311 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_add.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_add.cc
@@ -36,6 +36,7 @@ Status GatedAdd::ComputeInternal(OpKernelContext* context) const {
const Tensor* x = context->Input(0);
const Tensor* y = context->Input(1);
const Tensor* gate = context->Input(2);
+ const Tensor* gate_weight = context->Input(3);
const TensorShape& shape = x->Shape();
ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1");
@@ -46,13 +47,32 @@ Status GatedAdd::ComputeInternal(OpKernelContext* context) const {
const size_t last_axis = shape.NumDimensions() - 1;
const int64_t hidden_size = shape[last_axis];
ORT_RETURN_IF_NOT(hidden_size > 0, "X last dimension must be positive");
- ORT_RETURN_IF_NOT(gate->Shape()[last_axis] == 1, "gate last dimension must be 1");
for (size_t axis = 0; axis < last_axis; ++axis) {
ORT_RETURN_IF_NOT(gate->Shape()[axis] == shape[axis],
"gate dimension ", axis, " must match X");
}
Tensor* output = context->Output(0, shape);
+ if (gate_weight != nullptr) {
+ if constexpr (!std::is_same_v) {
+ return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "GatedAdd gate_weight currently supports BFloat16 only");
+ } else {
+ const int64_t input_width = gate->Shape()[last_axis];
+ ORT_RETURN_IF_NOT(input_width > 0 && gate_weight->Shape().NumDimensions() == 2 &&
+ gate_weight->Shape()[0] == input_width && gate_weight->Shape()[1] == 1,
+ "GatedAdd gate_weight must have shape [K, 1]");
+ return LaunchProjectedSigmoidGateMul(
+ Stream(context),
+ reinterpret_cast(y->Data()),
+ reinterpret_cast(gate->Data()),
+ reinterpret_cast(gate_weight->Data()),
+ reinterpret_cast(x->Data()),
+ reinterpret_cast(output->MutableData()),
+ shape.Size() / hidden_size, input_width, hidden_size);
+ }
+ }
+
+ ORT_RETURN_IF_NOT(gate->Shape()[last_axis] == 1, "gate last dimension must be 1");
return LaunchGatedAddKernel(
Stream(context),
reinterpret_cast(output->MutableData()),
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu
index 21e1ff6983a38..82aa1baf8d463 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu
@@ -4,7 +4,9 @@
#include
#include
#include
+#include
+#include
#include
#include "contrib_ops/cuda/bert/gated_add_impl.h"
@@ -17,6 +19,77 @@ namespace cuda {
namespace {
+constexpr int kProjectedGateThreads = 256;
+
+__device__ __forceinline__ float StableSigmoid(float value) {
+ if (value >= 0.0f) {
+ return 1.0f / (1.0f + expf(-value));
+ }
+ const float exp_value = expf(value);
+ return exp_value / (1.0f + exp_value);
+}
+
+__device__ __forceinline__ float RoundedBf16Sigmoid(float value) {
+ return __bfloat162float(__float2bfloat16_rn(StableSigmoid(value)));
+}
+
+__device__ float ProjectedRowGate(const __nv_bfloat16* activation, const __nv_bfloat16* weight,
+ int64_t row, int64_t input_width) {
+ using Reduce = cub::BlockReduce;
+ __shared__ typename Reduce::TempStorage reduction;
+ __shared__ float gate;
+ float dot = 0.0f;
+ for (int64_t k = threadIdx.x; k < input_width; k += kProjectedGateThreads) {
+ dot = fmaf(__bfloat162float(activation[row * input_width + k]),
+ __bfloat162float(weight[k]), dot);
+ }
+ const float total = Reduce(reduction).Sum(dot);
+ if (threadIdx.x == 0) {
+ const float value = __bfloat162float(__float2bfloat16_rn(total));
+ gate = RoundedBf16Sigmoid(value);
+ }
+ __syncthreads();
+ return gate;
+}
+
+__global__ void ProjectedSigmoidGateMulKernel(
+ const __nv_bfloat16* attention, const __nv_bfloat16* activation,
+ const __nv_bfloat16* weight, const __nv_bfloat16* residual, __nv_bfloat16* output,
+ int64_t input_width, int64_t output_width) {
+ const int64_t row = blockIdx.x;
+ const float gate = ProjectedRowGate(activation, weight, row, input_width);
+ const int64_t row_base = row * output_width;
+ if ((output_width & 7) == 0 &&
+ (reinterpret_cast(attention + row_base) & 15) == 0 &&
+ (reinterpret_cast(output + row_base) & 15) == 0 &&
+ (reinterpret_cast(residual + row_base) & 15) == 0) {
+ const int64_t vec_width = output_width >> 3;
+ const uint4* attn_vec = reinterpret_cast(attention + row_base);
+ const uint4* res_vec = reinterpret_cast(residual + row_base);
+ uint4* out_vec = reinterpret_cast(output + row_base);
+ for (int64_t v = threadIdx.x; v < vec_width; v += kProjectedGateThreads) {
+ const uint4 a_u4 = attn_vec[v];
+ const uint4 r_u4 = res_vec[v];
+ const __nv_bfloat16* a_bf = reinterpret_cast(&a_u4);
+ const __nv_bfloat16* r_bf = reinterpret_cast(&r_u4);
+ uint4 o_u4;
+ __nv_bfloat16* o_bf = reinterpret_cast<__nv_bfloat16*>(&o_u4);
+#pragma unroll
+ for (int m = 0; m < 8; ++m) {
+ const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(a_bf[m]) * gate);
+ o_bf[m] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(r_bf[m]));
+ }
+ out_vec[v] = o_u4;
+ }
+ return;
+ }
+ for (int64_t column = threadIdx.x; column < output_width; column += kProjectedGateThreads) {
+ const int64_t index = row_base + column;
+ const __nv_bfloat16 scaled = __float2bfloat16_rn(__bfloat162float(attention[index]) * gate);
+ output[index] = __float2bfloat16_rn(__bfloat162float(scaled) + __bfloat162float(residual[index]));
+ }
+}
+
template
__device__ __forceinline__ T RoundedMulAdd(T x, T y, T gate);
@@ -61,6 +134,30 @@ __global__ void GatedAddKernel(T* output,
}
}
+template
+__global__ void GatedAddVec8Kernel(uint4* output,
+ const uint4* x,
+ const uint4* y,
+ const T* gate,
+ int vec_count,
+ onnxruntime::cuda::fast_divmod vec_hidden_size) {
+ const int v = static_cast(blockIdx.x) * blockDim.x + threadIdx.x;
+ if (v < vec_count) {
+ const T g = gate[vec_hidden_size.div(v)];
+ const uint4 x_u4 = x[v];
+ const uint4 y_u4 = y[v];
+ uint4 out_u4;
+ const T* x_t = reinterpret_cast(&x_u4);
+ const T* y_t = reinterpret_cast(&y_u4);
+ T* out_t = reinterpret_cast(&out_u4);
+#pragma unroll
+ for (int m = 0; m < 8; ++m) {
+ out_t[m] = RoundedMulAdd(x_t[m], y_t[m], g);
+ }
+ output[v] = out_u4;
+ }
+}
+
} // namespace
template
@@ -76,6 +173,23 @@ Status LaunchGatedAddKernel(cudaStream_t stream,
}
constexpr int kThreads = 256;
+ if constexpr (sizeof(T) == 2) {
+ if ((hidden_size & 7) == 0 && (count >> 3) <= std::numeric_limits::max() &&
+ (reinterpret_cast(output) & 15) == 0 &&
+ (reinterpret_cast(x) & 15) == 0 &&
+ (reinterpret_cast(y) & 15) == 0) {
+ const int vec_count = static_cast(count >> 3);
+ const int vec_blocks = (vec_count - 1) / kThreads + 1;
+ GatedAddVec8Kernel<<>>(
+ reinterpret_cast(output),
+ reinterpret_cast(x),
+ reinterpret_cast(y),
+ gate,
+ vec_count,
+ onnxruntime::cuda::fast_divmod(static_cast(hidden_size >> 3)));
+ return CUDA_CALL(cudaGetLastError());
+ }
+ }
const int64_t blocks = (count - 1) / kThreads + 1;
ORT_RETURN_IF_NOT(blocks <= std::numeric_limits::max(),
"GatedAdd launch requires too many blocks");
@@ -100,6 +214,23 @@ INSTANTIATE_GATED_ADD(__nv_bfloat16)
#undef INSTANTIATE_GATED_ADD
+Status LaunchProjectedSigmoidGateMul(cudaStream_t stream, const uint16_t* attention,
+ const uint16_t* activation, const uint16_t* weight,
+ const uint16_t* residual, uint16_t* output,
+ int64_t rows, int64_t input_width, int64_t output_width) {
+ ORT_RETURN_IF_NOT(rows >= 0 && rows <= std::numeric_limits::max() &&
+ input_width > 0 && output_width > 0,
+ "GatedAdd projected gate dimensions exceed CUDA limits");
+ if (rows == 0) return Status::OK();
+ ProjectedSigmoidGateMulKernel<<(rows), kProjectedGateThreads, 0, stream>>>(
+ reinterpret_cast(attention),
+ reinterpret_cast(activation),
+ reinterpret_cast(weight),
+ reinterpret_cast(residual), reinterpret_cast<__nv_bfloat16*>(output),
+ input_width, output_width);
+ return CUDA_CALL(cudaGetLastError());
+}
+
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime
\ No newline at end of file
diff --git a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h
index 9df24bd846bb2..e1e2103789569 100644
--- a/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h
+++ b/onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h
@@ -20,6 +20,16 @@ Status LaunchGatedAddKernel(cudaStream_t stream,
int64_t count,
int64_t hidden_size);
+Status LaunchProjectedSigmoidGateMul(cudaStream_t stream,
+ const uint16_t* attention,
+ const uint16_t* activation,
+ const uint16_t* weight,
+ const uint16_t* residual,
+ uint16_t* output,
+ int64_t rows,
+ int64_t input_width,
+ int64_t output_width);
+
} // namespace cuda
} // namespace contrib
} // namespace onnxruntime
\ No newline at end of file
diff --git a/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc b/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc
index 1202fe55857a6..2e24fc8bcdc6c 100644
--- a/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc
+++ b/onnxruntime/contrib_ops/webgpu/bert/gated_add.cc
@@ -40,6 +40,10 @@ Status GatedAdd::ComputeInternal(ComputeContext& context) const {
const auto* x = context.Input(0);
const auto* y = context.Input(1);
const auto* gate = context.Input(2);
+ const auto* gate_weight = context.Input(3);
+ if (gate_weight != nullptr) {
+ return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, "GatedAdd gate_weight currently supports BFloat16 only");
+ }
const TensorShape& shape = x->Shape();
ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1");
diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
index b18836cf0b764..aca69d9276277 100644
--- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc
+++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc
@@ -6105,22 +6105,50 @@ ONNX_MS_OPERATOR_SET_SCHEMA(
}));
constexpr const char* GatedAdd_ver1_doc = R"DOC(
-Adds one tensor to another tensor scaled by a per-row gate:
+Adds one tensor to another tensor scaled by a per-row gate. X and Y have shape (..., C). The
+leading dimensions of gate must match the leading dimensions of X (there is no broadcasting over
+them); the optional gate_weight input selects how the per-row gate is obtained.
+
+Plain mode (gate_weight is not provided): gate has shape (..., 1) and is used directly as the
+gate, broadcast over C. No sigmoid is applied.
output = X + round_to_T(Y * gate)
-X and Y have shape (..., C), and gate has shape (..., 1). The gate is broadcast
-over C. For reduced-precision types, the product is rounded to T before the add,
-matching separate ONNX Mul and Add operators.
+For reduced-precision types, the product is rounded to T before the add, matching separate ONNX
+Mul and Add operators. This mode supports float, float16 and bfloat16, subject to execution provider kernel support.
+
+Projected mode (gate_weight is provided): gate holds an activation with shape (..., K) and
+gate_weight has shape (K, 1). For every row, the gate is computed by projecting the activation
+with gate_weight and applying a sigmoid:
+
+ p = round_to_bfloat16(sum_k(gate[..., k] * gate_weight[k]))
+ s = round_to_bfloat16(sigmoid(p))
+ output = round_to_bfloat16(X + round_to_bfloat16(Y * s))
+
+The dot product, the sigmoid and the multiplication and addition are evaluated in FP32 (the
+sigmoid is always computed in FP32, and the order of the FP32 accumulation over K is not
+specified); the result is rounded to bfloat16 after the projection, after the sigmoid, after the
+multiplication and after the addition. This mode is only implemented by the CUDA execution
+provider for bfloat16 tensors. Other execution providers and types return NOT_IMPLEMENTED when
+gate_weight is provided.
)DOC";
ONNX_MS_OPERATOR_SET_SCHEMA(
GatedAdd, 1,
OpSchema()
.SetDoc(GatedAdd_ver1_doc)
- .Input(0, "X", "Unscaled input with shape (..., C).", "T")
- .Input(1, "Y", "Input scaled by gate, with the same shape as X.", "T")
- .Input(2, "gate", "Per-row gate with shape (..., 1).", "T")
+ .Input(0, "X",
+ "Unscaled input with shape (..., C).",
+ "T")
+ .Input(1, "Y",
+ "Input scaled by the gate, with the same shape as X.",
+ "T")
+ .Input(2, "gate",
+ "In plain mode, the per-row gate with shape (..., 1). When gate_weight is provided, the activation with shape (..., K) that is projected with gate_weight and passed through a sigmoid to produce the per-row gate. Its leading dimensions must match X.",
+ "T")
+ .Input(3, "gate_weight",
+ "Optional projection weight with shape (K, 1), where K is the last dimension of gate. When present, the gate is sigmoid(gate . gate_weight) computed per row in FP32 (projected mode). Only supported by the CUDA execution provider with bfloat16.",
+ "T", OpSchema::Optional)
.Output(0, "output", "Gated sum with the same shape as X.", "T")
.TypeConstraint("T",
{"tensor(float)", "tensor(float16)", "tensor(bfloat16)"},
diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc
index 67bf96d8f3199..1c1ca9088b7a5 100755
--- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc
+++ b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc
@@ -1392,9 +1392,11 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kO
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, ArgMax);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, double, ArgMax);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, ArgMax);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ArgMax);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, float, ArgMin);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, double, ArgMin);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, MLFloat16, ArgMin);
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 13, BFloat16, ArgMin);
// OpSet 14
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCudaExecutionProvider, kOnnxDomain, 14, CumSum);
@@ -2722,9 +2724,11 @@ static Status RegisterCudaKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
BuildKernelCreateInfo,
+ BuildKernelCreateInfo,
// OpSet 14
BuildKernelCreateInfo,
diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
index 60ffd772c97f1..5fc35f5e68858 100644
--- a/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
+++ b/onnxruntime/core/providers/cuda/reduction/reduction_functions.cu
@@ -777,6 +777,12 @@ struct ArgReductionIdentity {
__device__ __forceinline__ static half Highest() { return __ushort_as_half(static_cast(0x7C00U)); }
};
+template <>
+struct ArgReductionIdentity {
+ __device__ __forceinline__ static BFloat16 Lowest() { return BFloat16::FromBits(static_cast(0xFF80U)); }
+ __device__ __forceinline__ static BFloat16 Highest() { return BFloat16::FromBits(static_cast(0x7F80U)); }
+};
+
template
__device__ __forceinline__ T arg_reduction_identity() {
return IsArgMax ? ArgReductionIdentity::Lowest() : ArgReductionIdentity::Highest();
@@ -1163,6 +1169,7 @@ Status arg_min_max_last_axis(cudaStream_t stream, const TIn* input, int64_t* out
int64_t* output, int m, int n, void* buffer, \
size_t buffer_size)
INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(half);
+INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(BFloat16);
INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(float);
INSTANTIATE_ARG_MIN_MAX_LAST_AXIS(double);
#undef INSTANTIATE_ARG_MIN_MAX_LAST_AXIS
diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
index 3d919ac04ecbf..18cc5fdee0668 100644
--- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
+++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc
@@ -527,9 +527,10 @@ Status ReduceComputeCore(const AllocatorPtr& gpu_allocator, const CudaKernel* ke
cudnnDataType_t cudnn_type_X = CUDNN_DATA_FLOAT;
// Reducesum with BFP16 is not supported by cudnn, so convert input to fp32 then call cudnn
- if ((ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES && std::is_same::value) ||
+ if ((ReduceTensorIndices == CUDNN_REDUCE_TENSOR_FLATTENED_INDICES &&
+ (std::is_same::value || std::is_same::value)) ||
(ReduceTensorIndices == CUDNN_REDUCE_TENSOR_NO_INDICES && std::is_same::value)) {
- // ArgMax/ArgMin with FP16 are not supported by cudnn, so convert input to fp32 then call cudnn
+ // ArgMax/ArgMin with FP16/BF16 are not supported by cudnn, so convert input to fp32 then call cudnn
temp_X = AllocateScratchBuffer(gpu_allocator, kernel, input_count, compute_stream);
Impl_Cast(stream, reinterpret_cast(input.Data()), temp_X.get(), input_shape.Size());
} else {
@@ -1031,10 +1032,12 @@ template std::unique_ptr ReduceCompute({1, 4, 2048});
}
+TEST(ContribOpGatedAddTest, Float16NonVectorizedHiddenSize) {
+ RunGatedAddTest({3, 5, 7});
+}
+
+TEST(ContribOpGatedAddTest, BFloat16NonVectorizedHiddenSize) {
+ if (!CudaHasBF16Support()) {
+ GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later";
+ }
+ RunGatedAddTest({3, 5, 7});
+}
+
TEST(ContribOpGatedAddTest, MismatchedYShape) {
auto execution_providers = AvailableGatedOpExecutionProviders();
for (auto& ep : execution_providers) {
@@ -179,5 +190,189 @@ TEST(ContribOpGatedAddTest, MismatchedYShape) {
}
}
+TEST(ContribOpGatedAddTest, ProjectedGateBFloat16) {
+ if (!CudaHasBF16Support()) {
+ GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later";
+ }
+ struct ProjectionCase {
+ int64_t width;
+ int64_t input_width;
+ };
+ for (const ProjectionCase& test_case : {ProjectionCase{8, 16}, ProjectionCase{2560, 16}, ProjectionCase{2560, 2560},
+ ProjectionCase{259, 300}}) {
+ constexpr int64_t kRows = 4;
+ const int64_t width = test_case.width;
+ const int64_t input_width = test_case.input_width;
+ std::vector x(kRows * width);
+ std::vector y(kRows * width);
+ std::vector act(kRows * input_width);
+ std::vector weight(input_width);
+ std::vector expected(kRows * width);
+ for (int64_t k = 0; k < input_width; ++k) {
+ weight[k] = BFloat16(static_cast((k % 5) - 2) * 0.25f);
+ }
+ for (int64_t r = 0; r < kRows; ++r) {
+ float dot = 0.0f;
+ for (int64_t k = 0; k < input_width; ++k) {
+ act[r * input_width + k] = BFloat16(static_cast(((r + k) % 7) - 3) * 0.125f);
+ dot = std::fma(act[r * input_width + k].ToFloat(), weight[k].ToFloat(), dot);
+ }
+ const float gate = BFloat16(dot).ToFloat();
+ const float exp_value = std::exp(gate >= 0.0f ? -gate : gate);
+ const float sigmoid = gate >= 0.0f ? 1.0f / (1.0f + exp_value) : exp_value / (1.0f + exp_value);
+ const BFloat16 rounded_sigmoid(sigmoid);
+ for (int64_t c = 0; c < width; ++c) {
+ const int64_t idx = r * width + c;
+ x[idx] = BFloat16(static_cast((idx % 17) - 8) * 0.13f);
+ y[idx] = BFloat16(static_cast((idx % 11) - 5) * 0.37f);
+ const BFloat16 product(y[idx].ToFloat() * rounded_sigmoid.ToFloat());
+ expected[idx] = BFloat16(x[idx].ToFloat() + product.ToFloat());
+ }
+ }
+ std::vector> providers;
+ if (auto cuda_ep = DefaultCudaExecutionProvider()) {
+ providers.push_back(std::move(cuda_ep));
+ }
+ if (providers.empty()) GTEST_SKIP();
+ OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain);
+ tester.AddInput("X", {1, kRows, width}, x);
+ tester.AddInput("Y", {1, kRows, width}, y);
+ tester.AddInput("gate", {1, kRows, input_width}, act);
+ tester.AddInput("gate_weight", {input_width, 1}, weight);
+ tester.AddOutput("output", {1, kRows, width}, expected);
+ tester.SetOutputTolerance(0.0f, 0.0f);
+ tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+ }
+}
+
+TEST(ContribOpGatedAddTest, ProjectedGateK1Regression) {
+ for (auto& ep : AvailableGatedOpExecutionProviders()) {
+ SCOPED_TRACE("EP (float): " + ep->Type());
+ OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain);
+ tester.AddInput("X", {1, 1}, {0.0f});
+ tester.AddInput("Y", {1, 1}, {2.0f});
+ tester.AddInput("gate", {1, 1}, {0.0f});
+ tester.AddInput("gate_weight", {1, 1}, {1.0f});
+ tester.AddOutput("output", {1, 1}, {1.0f});
+
+ std::vector> providers;
+ providers.push_back(std::move(ep));
+ tester.Run(OpTester::ExpectResult::kExpectFailure, "GatedAdd gate_weight currently supports BFloat16 only",
+ {}, nullptr, &providers);
+ }
+
+ for (auto& ep : AvailableGatedOpExecutionProviders()) {
+ SCOPED_TRACE("EP (MLFloat16): " + ep->Type());
+ OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain);
+ tester.AddInput("X", {1, 1}, ToTensorType({0.0f}));
+ tester.AddInput("Y", {1, 1}, ToTensorType({2.0f}));
+ tester.AddInput("gate", {1, 1}, ToTensorType({0.0f}));
+ tester.AddInput("gate_weight", {1, 1}, ToTensorType({1.0f}));
+ tester.AddOutput("output", {1, 1}, ToTensorType({1.0f}));
+
+ std::vector> providers;
+ providers.push_back(std::move(ep));
+ tester.Run(OpTester::ExpectResult::kExpectFailure, "GatedAdd gate_weight currently supports BFloat16 only",
+ {}, nullptr, &providers);
+ }
+
+ if (CudaHasBF16Support()) {
+ std::vector> providers;
+ if (auto cuda_ep = DefaultCudaExecutionProvider()) {
+ providers.push_back(std::move(cuda_ep));
+ }
+ if (!providers.empty()) {
+ OpTester tester("GatedAdd", 1, onnxruntime::kMSDomain);
+ tester.AddInput("X", {1, 1}, ToTensorType({0.0f}));
+ tester.AddInput("Y", {1, 1}, ToTensorType({2.0f}));
+ tester.AddInput("gate", {1, 1}, ToTensorType({0.0f}));
+ tester.AddInput("gate_weight", {1, 1}, ToTensorType({1.0f}));
+ tester.AddOutput("output", {1, 1}, ToTensorType({1.0f}));
+ tester.SetOutputTolerance(0.0f, 0.0f);
+ tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+ }
+ }
+}
+
+TEST(ContribOpGatedAddTest, ArgMaxBFloat16) {
+ if (!CudaHasBF16Support()) {
+ GTEST_SKIP() << "bfloat16 requires compute capability 8.0 or later";
+ }
+ for (bool keepdims : {false, true}) {
+ for (const char* op : {"ArgMax", "ArgMin"}) {
+ const bool is_max = std::string(op) == "ArgMax";
+ std::vector> providers;
+ if (auto cuda_ep = DefaultCudaExecutionProvider()) {
+ providers.push_back(std::move(cuda_ep));
+ }
+ if (providers.empty()) GTEST_SKIP();
+ OpTester test(op, 13);
+ test.AddAttribute("axis", -1);
+ test.AddAttribute("keepdims", keepdims ? 1 : 0);
+ test.AddInput("logits", {2, 4},
+ ToTensorType({1.0f, 5.0f, 5.0f, -1.0f, 4.0f, 3.0f, 2.0f, 1.0f}));
+ test.AddOutput("indices", keepdims ? std::vector{2, 1} : std::vector{2},
+ is_max ? std::vector{1, 0} : std::vector{3, 3});
+ test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+ }
+
+ for (int64_t cols : {int64_t{1024}, int64_t{40000}}) {
+ constexpr int64_t kRows = 2;
+ const std::vector max_indices{17, cols - 5};
+ const std::vector min_indices{29, cols / 2};
+ std::vector logits(static_cast(kRows * cols));
+ for (int64_t r = 0; r < kRows; ++r) {
+ for (int64_t c = 0; c < cols; ++c) {
+ logits[static_cast(r * cols + c)] = static_cast((c + r) % 31) * 0.5f;
+ }
+ logits[static_cast(r * cols + max_indices[static_cast(r)])] = 256.0f;
+ logits[static_cast(r * cols + cols - 1)] = 256.0f;
+ logits[static_cast(r * cols + min_indices[static_cast(r)])] = -256.0f;
+ logits[static_cast(r * cols + cols - 2)] = -256.0f;
+ }
+ for (const char* op : {"ArgMax", "ArgMin"}) {
+ const bool is_max = std::string(op) == "ArgMax";
+ std::vector> providers;
+ if (auto cuda_ep = DefaultCudaExecutionProvider()) {
+ providers.push_back(std::move(cuda_ep));
+ }
+ if (providers.empty()) GTEST_SKIP();
+ OpTester test(op, 13);
+ test.AddAttribute("axis", -1);
+ test.AddAttribute("keepdims", keepdims ? 1 : 0);
+ test.AddInput("logits", {kRows, cols}, ToTensorType(logits));
+ test.AddOutput("indices", keepdims ? std::vector{kRows, 1} : std::vector{kRows},
+ is_max ? max_indices : min_indices);
+ test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+ }
+ }
+
+ for (const char* op : {"ArgMax", "ArgMin"}) {
+ const bool is_max = std::string(op) == "ArgMax";
+ std::vector> providers;
+ if (auto cuda_ep = DefaultCudaExecutionProvider()) {
+ providers.push_back(std::move(cuda_ep));
+ }
+ if (providers.empty()) GTEST_SKIP();
+ OpTester test(op, 13);
+ test.AddAttribute("axis", 1);
+ test.AddAttribute("keepdims", keepdims ? 1 : 0);
+ test.AddInput("logits", {3, 2, 2},
+ ToTensorType({1.0f, 4.0f,
+ 3.0f, 2.0f,
+
+ 8.0f, 5.0f,
+ 7.0f, 6.0f,
+
+ 9.0f, 12.0f,
+ 11.0f, 10.0f}));
+ test.AddOutput("indices", keepdims ? std::vector{3, 1, 2} : std::vector{3, 2},
+ is_max ? std::vector{1, 0, 0, 1, 1, 0}
+ : std::vector{0, 1, 1, 0, 0, 1});
+ test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &providers);
+ }
+ }
+}
+
} // namespace test
} // namespace onnxruntime
diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
index 03a582af013a5..228a11e633107 100644
--- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
+++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
@@ -6983,6 +6983,54 @@ TEST(ReductionOpTest, ReduceMin_float_Opset20_NoKeepdims_Cuda) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
}
+TEST(ReductionOpTest, ArgMax_bfloat16_non_last_axis_Cuda) {
+ for (int64_t keepdims : {int64_t{0}, int64_t{1}}) {
+ OpTester test("ArgMax", 13);
+ test.AddAttribute("axis", int64_t{1});
+ test.AddAttribute("keepdims", keepdims);
+ test.AddInput("data", {3, 2, 2},
+ MakeBFloat16({1.0f, 4.0f,
+ 3.0f, 2.0f,
+
+ 8.0f, 5.0f,
+ 7.0f, 6.0f,
+
+ 9.0f, 12.0f,
+ 11.0f, 10.0f}));
+ test.AddOutput("reduced", keepdims ? std::vector{3, 1, 2} : std::vector{3, 2},
+ {1, 0,
+ 0, 1,
+ 1, 0});
+ std::vector> execution_providers;
+ execution_providers.push_back(DefaultCudaExecutionProvider());
+ test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+ }
+}
+
+TEST(ReductionOpTest, ArgMin_bfloat16_non_last_axis_Cuda) {
+ for (int64_t keepdims : {int64_t{0}, int64_t{1}}) {
+ OpTester test("ArgMin", 13);
+ test.AddAttribute("axis", int64_t{1});
+ test.AddAttribute("keepdims", keepdims);
+ test.AddInput("data", {3, 2, 2},
+ MakeBFloat16({1.0f, 4.0f,
+ 3.0f, 2.0f,
+
+ 8.0f, 5.0f,
+ 7.0f, 6.0f,
+
+ 9.0f, 12.0f,
+ 11.0f, 10.0f}));
+ test.AddOutput("reduced", keepdims ? std::vector{3, 1, 2} : std::vector{3, 2},
+ {0, 1,
+ 1, 0,
+ 0, 1});
+ std::vector> execution_providers;
+ execution_providers.push_back(DefaultCudaExecutionProvider());
+ test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+ }
+}
+
#endif // defined(USE_CUDA)
// =============================================================================
diff --git a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc
index d4e5d5e923957..2f4af0288f38b 100644
--- a/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc
+++ b/onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc
@@ -421,6 +421,12 @@ struct ArgTestTraits {
static double ToDouble(half value) { return static_cast(__half2float(value)); }
};
+template <>
+struct ArgTestTraits {
+ static BFloat16 FromFloat(float value) { return BFloat16(value); }
+ static double ToDouble(BFloat16 value) { return static_cast(value.ToFloat()); }
+};
+
// Mirrors the sequential semantics of the reduction: seed with the first
// element, keep the first occurrence of the extreme value and never select a
// NaN (every comparison against NaN is false, so a leading NaN wins the row).
@@ -584,6 +590,7 @@ TEST(ReductionFunctionsTest, ArgMinMaxLastAxisManyRows) {
TEST(ReductionFunctionsTest, ArgMinMaxLastAxisTypes) {
for (int n : {64, 1024, 65536}) {
TestArgMinMaxLastAxisRandom(3, n, n);
+ TestArgMinMaxLastAxisRandom(3, n, n);
TestArgMinMaxLastAxisRandom(3, n, n);
TestArgMinMaxLastAxisRandom(3, n, n);
}
@@ -593,6 +600,7 @@ TEST(ReductionFunctionsTest, ArgMinMaxLastAxisSpecialValues) {
// 64 -> one thread per row, 1024 -> single block per row, 200000 -> many blocks per row.
for (int n : {64, 1024, 200000}) {
TestArgMinMaxLastAxisSpecialValues(n);
+ TestArgMinMaxLastAxisSpecialValues(n);
TestArgMinMaxLastAxisSpecialValues(n);
TestArgMinMaxLastAxisSpecialValues(n);
}