[CUDA] Add BFloat16 ArgMax/ArgMin support and extend GatedAdd with projected gate

docs/ContribOperators.md

@@ -2387,27 +2387,48 @@ This version of the operator has been available since version 1 of the 'com.micr
 
 ### <a name="com.microsoft.GatedAdd"></a><a name="com.microsoft.gatedadd">**com.microsoft.GatedAdd**</a>
 
-  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)
 
 <dl>
 <dt><tt>X</tt> : T</dt>
 <dd>Unscaled input with shape (..., C).</dd>
 <dt><tt>Y</tt> : T</dt>
-<dd>Input scaled by gate, with the same shape as X.</dd>
+<dd>Input scaled by the gate, with the same shape as X.</dd>
 <dt><tt>gate</tt> : T</dt>
-<dd>Per-row gate with shape (..., 1).</dd>
+<dd>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.</dd>
+<dt><tt>gate_weight</tt> (optional) : T</dt>
+<dd>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.</dd>
 </dl>
 
 #### Outputs

docs/OperatorKernels.md

@@ -600,7 +600,7 @@ The **OpSet Version** column uses the following notation:
 |FusedConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *in* Z:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |FusedGemm|*in* A:**T**<br> *in* B:**T**<br> *in* C:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |FusedMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float)|
-|GatedAdd|*in* X:**T**<br> *in* Y:**T**<br> *in* gate:**T**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
+|GatedAdd|*in* X:**T**<br> *in* Y:**T**<br> *in* gate:**T**<br> *in* gate_weight:**T**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|
 |GatedRMSNorm|*in* X:**T**<br> *in* scale:**T**<br> *in* gate:**T**<br> *out* Y:**T**|1+|**T** = tensor(float), tensor(float16)|
 |GatherBlockQuantized|*in* data:**T1**<br> *in* indices:**Tind**<br> *in* scales:**T2**<br> *in* zero_points:**T1**<br> *out* output:**T2**|1+|**T1** = tensor(float4e2m1), tensor(float8e4m3fn), tensor(float8e4m3fnuz), tensor(float8e5m2), tensor(float8e5m2fnuz), tensor(int4), tensor(uint4), tensor(uint8)<br/> **T2** = tensor(float), tensor(float16)<br/> **Tind** = tensor(int32), tensor(int64)|
 |GatherND|*in* data:**T**<br> *in* indices:**Tind**<br> *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)<br/> **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**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
 |And|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T1**|7+|**T** = tensor(bool)<br/> **T1** = tensor(bool)|
-|ArgMax|*in* data:**T**<br> *out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(float16)|
+|ArgMax|*in* data:**T**<br> *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**<br> *out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(float16)|
+|ArgMin|*in* data:**T**<br> *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**<br> *in* K:**T1**<br> *in* V:**T2**<br> *in* attn_mask:**U**<br> *in* past_key:**T1**<br> *in* past_value:**T2**<br> *in* nonpad_kv_seqlen:**tensor(int64)**<br> *out* Y:**T1**<br> *out* present_key:**T1**<br> *out* present_value:**T2**<br> *out* qk_matmul_output:**T1**<br><br>or<br><br>*in* Q:**T1**<br> *in* K:**T1**<br> *in* V:**T2**<br> *in* attn_mask:**U**<br> *in* past_key:**T1**<br> *in* past_value:**T2**<br> *out* Y:**T1**<br> *out* present_key:**T1**<br> *out* present_value:**T2**<br> *out* qk_matmul_output:**T1**|24+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **T2** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **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**<br> *in* bias:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
 |FusedConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *in* Z:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |FusedMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
-|GatedAdd|*in* X:**T**<br> *in* Y:**T**<br> *in* gate:**T**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
+|GatedAdd|*in* X:**T**<br> *in* Y:**T**<br> *in* gate:**T**<br> *in* gate_weight:**T**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
 |GatedDeltaNet|*in* query:**T**<br> *in* key:**T**<br> *in* value:**T**<br> *in* cu_seqlens:**TI**<br> *in* decay:**TS**<br> *in* beta:**TS**<br> *in* initial_state:**TS**<br> *in* a_log:**TS**<br> *in* dt_bias:**TS**<br> *in* capture_count:**TI**<br> *in* state_update_active:**TI**<br> *out* output:**T**<br> *out* final_state:**TS**<br> *out* state_update:**TS**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)<br/> **TI** = tensor(int32)<br/> **TS** = tensor(float)|
 |GatedRMSNorm|*in* X:**T**<br> *in* scale:**T**<br> *in* gate:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float), tensor(float16)|
 |GatedRelativePositionBias|*in* query_layer:**T**<br> *in* query_bias:**T**<br> *in* rel_pos:**T**<br> *in* weight:**T**<br> *in* bias:**T**<br> *in* eco_a:**T**<br> *in* token_offset:**M**<br> *out* output:**T**|1+|**T** = tensor(float), tensor(float16)|

onnxruntime/contrib_ops/cpu/bert/gated_add.cc

@@ -46,6 +46,10 @@ Status GatedAdd<T>::Compute(OpKernelContext* context) const {
   const Tensor* x = context->Input<Tensor>(0);
   const Tensor* y = context->Input<Tensor>(1);
   const Tensor* gate = context->Input<Tensor>(2);
+  const Tensor* gate_weight = context->Input<Tensor>(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");

onnxruntime/contrib_ops/cuda/bert/gated_add.cc

@@ -36,6 +36,7 @@ Status GatedAdd<T>::ComputeInternal(OpKernelContext* context) const {
   const Tensor* x = context->Input<Tensor>(0);
   const Tensor* y = context->Input<Tensor>(1);
   const Tensor* gate = context->Input<Tensor>(2);
+  const Tensor* gate_weight = context->Input<Tensor>(3);
   const TensorShape& shape = x->Shape();
 
   ORT_RETURN_IF_NOT(shape.NumDimensions() >= 1, "X must have rank >= 1");
@@ -46,13 +47,32 @@ Status GatedAdd<T>::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<T, BFloat16>) {
+      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<const uint16_t*>(y->Data<BFloat16>()),
+          reinterpret_cast<const uint16_t*>(gate->Data<BFloat16>()),
+          reinterpret_cast<const uint16_t*>(gate_weight->Data<BFloat16>()),
+          reinterpret_cast<const uint16_t*>(x->Data<BFloat16>()),
+          reinterpret_cast<uint16_t*>(output->MutableData<BFloat16>()),
+          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<CudaT>(
       Stream(context),
       reinterpret_cast<CudaT*>(output->MutableData<T>()),

onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu

@@ -4,7 +4,9 @@
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
+#include <cub/block/block_reduce.cuh>
 
+#include <cmath>
 #include <limits>
 
 #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<float, kProjectedGateThreads>;
+  __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<uintptr_t>(attention + row_base) & 15) == 0 &&
+      (reinterpret_cast<uintptr_t>(output + row_base) & 15) == 0 &&
+      (reinterpret_cast<uintptr_t>(residual + row_base) & 15) == 0) {
+    const int64_t vec_width = output_width >> 3;
+    const uint4* attn_vec = reinterpret_cast<const uint4*>(attention + row_base);
+    const uint4* res_vec = reinterpret_cast<const uint4*>(residual + row_base);
+    uint4* out_vec = reinterpret_cast<uint4*>(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<const __nv_bfloat16*>(&a_u4);
+      const __nv_bfloat16* r_bf = reinterpret_cast<const __nv_bfloat16*>(&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 <typename T>
 __device__ __forceinline__ T RoundedMulAdd(T x, T y, T gate);
 
@@ -61,6 +134,30 @@ __global__ void GatedAddKernel(T* output,
   }
 }
 
+template <typename T>
+__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<int>(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<const T*>(&x_u4);
+    const T* y_t = reinterpret_cast<const T*>(&y_u4);
+    T* out_t = reinterpret_cast<T*>(&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 <typename T>
@@ -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<int>::max() &&
+        (reinterpret_cast<uintptr_t>(output) & 15) == 0 &&
+        (reinterpret_cast<uintptr_t>(x) & 15) == 0 &&
+        (reinterpret_cast<uintptr_t>(y) & 15) == 0) {
+      const int vec_count = static_cast<int>(count >> 3);
+      const int vec_blocks = (vec_count - 1) / kThreads + 1;
+      GatedAddVec8Kernel<T><<<vec_blocks, kThreads, 0, stream>>>(
+          reinterpret_cast<uint4*>(output),
+          reinterpret_cast<const uint4*>(x),
+          reinterpret_cast<const uint4*>(y),
+          gate,
+          vec_count,
+          onnxruntime::cuda::fast_divmod(static_cast<int>(hidden_size >> 3)));
+      return CUDA_CALL(cudaGetLastError());
+    }
+  }
   const int64_t blocks = (count - 1) / kThreads + 1;
   ORT_RETURN_IF_NOT(blocks <= std::numeric_limits<int>::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<int>::max() &&
+                        input_width > 0 && output_width > 0,
+                    "GatedAdd projected gate dimensions exceed CUDA limits");
+  if (rows == 0) return Status::OK();
+  ProjectedSigmoidGateMulKernel<<<static_cast<int>(rows), kProjectedGateThreads, 0, stream>>>(
+      reinterpret_cast<const __nv_bfloat16*>(attention),
+      reinterpret_cast<const __nv_bfloat16*>(activation),
+      reinterpret_cast<const __nv_bfloat16*>(weight),
+      reinterpret_cast<const __nv_bfloat16*>(residual), reinterpret_cast<__nv_bfloat16*>(output),
+      input_width, output_width);
+  return CUDA_CALL(cudaGetLastError());
+}
+
 }  // namespace cuda
 }  // namespace contrib
 }  // namespace onnxruntime

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

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");