Update MoE and qMoE spec

docs/ContribOperators.md

@@ -3106,15 +3106,15 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>router_probs</tt> : T</dt>
 <dd>2D input tensor with shape (num_rows, num_experts)</dd>
 <dt><tt>fc1_experts_weights</tt> : T</dt>
-<dd>3D input tensor with shape (num_experts, hidden_size, inter_size), or (num_experts, hidden_size, 2 * inter_size) for swiglu</dd>
+<dd>3D input tensor with shape (num_experts, inter_size, hidden_size), or (num_experts, 2 * inter_size, hidden_size) for swiglu</dd>
 <dt><tt>fc1_experts_bias</tt> (optional) : T</dt>
 <dd>2D optional input tensor with shape (num_experts, inter_size), or (num_experts, 2 * inter_size) for swiglu</dd>
 <dt><tt>fc2_experts_weights</tt> : T</dt>
-<dd>3D input tensor with shape (num_experts, inter_size, hidden_size)</dd>
+<dd>3D input tensor with shape (num_experts, hidden_size, inter_size)</dd>
 <dt><tt>fc2_experts_bias</tt> (optional) : T</dt>
 <dd>2D optional input tensor with shape (num_experts, hidden_size)</dd>
 <dt><tt>fc3_experts_weights</tt> (optional) : T</dt>
-<dd>3D optional input tensor with shape (num_experts, hidden_size, inter_size)</dd>
+<dd>3D optional input tensor with shape (num_experts, inter_size, hidden_size)</dd>
 <dt><tt>fc3_experts_bias</tt> (optional) : T</dt>
 <dd>2D optional input tensor with shape (num_experts, inter_size)</dd>
 </dl>
@@ -4542,19 +4542,19 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dt><tt>router_probs</tt> : T</dt>
 <dd>2D input tensor with shape (num_rows, num_experts)</dd>
 <dt><tt>fc1_experts_weights</tt> : T1</dt>
-<dd>3D input tensor with shape (num_experts, hidden_size, inter_size) or (num_experts, hidden_size, inter_size / 2). For swiglu, shape can be (num_experts, hidden_size, 2 * inter_size) or (num_experts, hidden_size, inter_size).</dd>
+<dd>3D input tensor with shape (num_experts, inter_size, hidden_size), or (num_experts, inter_size, hidden_size / 2) for 4 bits. For swiglu, shape can be (num_experts, 2 * inter_size, hidden_size), or (num_experts, 2 * inter_size, hidden_size / 2) for 4 bits.</dd>
 <dt><tt>fc1_scales</tt> : T2</dt>
 <dd>2D input tensor with shape (num_experts, inter_size), or (num_experts, 2 * inter_size) for swiglu</dd>
 <dt><tt>fc1_experts_bias</tt> (optional) : T</dt>
 <dd>2D optional input tensor with shape (num_experts, inter_size), or (num_experts, 2 * inter_size) for swiglu</dd>
 <dt><tt>fc2_experts_weights</tt> : T1</dt>
-<dd>3D input tensor with shape (num_experts, inter_size, hidden_size) or (num_experts, inter_size, hidden_size / 2)</dd>
+<dd>3D input tensor with shape (num_experts, hidden_size, inter_size) or (num_experts, hidden_size, inter_size / 2) for 4 bits</dd>
 <dt><tt>fc2_scales</tt> : T2</dt>
 <dd>2D input tensor with shape (num_experts, hidden_size)</dd>
 <dt><tt>fc2_experts_bias</tt> (optional) : T</dt>
 <dd>2D optional input tensor with shape (num_experts, hidden_size)</dd>
 <dt><tt>fc3_experts_weights</tt> (optional) : T1</dt>
-<dd>3D optional input tensor with shape (num_experts, hidden_size, inter_size) or (num_experts, hidden_size, inter_size / 2)</dd>
+<dd>3D optional input tensor with shape (num_experts, inter_size, hidden_size) or (num_experts, inter_size, hidden_size / 2)</dd>
 <dt><tt>fc3_scales</tt> (optional) : T2</dt>
 <dd>2D optional input tensor with shape (num_experts, inter_size)</dd>
 <dt><tt>fc3_experts_bias</tt> (optional) : T</dt>
@@ -4575,7 +4575,7 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>Constrain input and output types to float tensors.</dd>
 <dt><tt>T1</tt> : tensor(uint8)</dt>
 <dd>Constrain weights type to uint8 tensors.</dd>
-<dt><tt>T2</tt> : tensor(float), tensor(float16)</dt>
+<dt><tt>T2</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
 <dd>Constrain scales type to float tensors.</dd>
 </dl>
 

onnxruntime/contrib_ops/cpu/quantization/moe_helper.h

@@ -0,0 +1,128 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/common/common.h"
+#include "core/providers/common.h"
+#include "core/framework/tensor_shape.h"
+#include "core/util/shape_checker.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+enum class MoEParallelType {
+  None = 0,
+  EP = 1,
+  TP = 2,
+  EPAndTP = 3,
+};
+
+struct MoEParameters {
+  MoEParameters() {}
+  explicit MoEParameters(int64_t tensor_shards) : tensor_shards(tensor_shards) {}
+  int64_t num_rows;
+  int64_t num_experts;
+  int64_t local_num_experts;
+  int64_t hidden_size;
+  int64_t inter_size;
+
+  MoEParallelType parallel_type;
+  int64_t tensor_shards{1};
+};
+
+namespace moe_helper {
+
+template <typename Tensor>
+Status CheckInputs(MoEParameters& parameters,
+                   const Tensor* input,                // required
+                   const Tensor* router_probs,         // required
+                   const Tensor* fc1_experts_weights,  // required
+                   const Tensor* fc1_experts_bias,     // optional
+                   const Tensor* fc1_experts_scales,   // required for qMoE; NULL for MOE
+                   const Tensor* fc2_experts_weights,  // required
+                   const Tensor* fc2_experts_bias,     // optional
+                   const Tensor* fc2_experts_scales,   // required for qMoE; NULL for MOE
+                   const Tensor* fc3_experts_weights,  // optional
+                   const Tensor* fc3_experts_bias,     // optional
+                   const Tensor* fc3_experts_scales,   // required for qMoE; NULL for MOE
+                   const int pack_size,                // number of weights packed together (like 2 for uint4 packed to uint8)
+                   const bool is_fused_swiglu) {
+
+  // Check dimensions of input to avoid input_dims index out of range. CHECK_TENSOR_SHAPE will verify each tensor later.
+  ASSERT_TENSOR_2D_OR_3D(input);
+  ASSERT_TENSOR_3D(fc1_experts_weights);
+  ASSERT_TENSOR_3D(fc2_experts_weights);
+  ASSERT_TENSOR_2D(router_probs);
+
+  const auto& input_dims = input->Shape().GetDims();
+  const auto& router_probs_dims = router_probs->Shape().GetDims();
+  const auto& fc1_experts_weights_dims = fc1_experts_weights->Shape().GetDims();
+  const auto& fc2_experts_weights_dims = fc2_experts_weights->Shape().GetDims();
+
+  int64_t num_rows = input_dims.size() == 2 ? input_dims[0] : input_dims[0] * input_dims[1];
+  int64_t hidden_size = input_dims[input_dims.size() - 1];
+  int64_t local_num_experts = fc1_experts_weights_dims[0];
+  int64_t num_experts = router_probs_dims[1];
+  int64_t inter_size = (fc2_experts_weights_dims[1] * fc2_experts_weights_dims[2] * pack_size) / hidden_size;
+  const bool legacy_shape = hidden_size != inter_size && fc2_experts_weights_dims[1] == inter_size;
+
+  // Fused swiglu doubles the output dimension of FC1 since it fused two GEMMs into one.
+  const int fc1_inter_size = is_fused_swiglu ? 2 * inter_size : inter_size;
+
+  if (legacy_shape) {
+    // legacy shape does not match the memory layout. This is for backward compatible
+    CHECK_TENSOR_SHAPE(fc1_experts_weights, num_experts, hidden_size, fc1_inter_size / pack_size);
+    CHECK_TENSOR_SHAPE(fc2_experts_weights, num_experts, inter_size, hidden_size / pack_size);
+    CHECK_TENSOR_SHAPE(fc3_experts_weights, num_experts, hidden_size, inter_size / pack_size);
+  } else {
+    CHECK_TENSOR_SHAPE(fc1_experts_weights, num_experts, fc1_inter_size, hidden_size / pack_size);
+    CHECK_TENSOR_SHAPE(fc2_experts_weights, num_experts, hidden_size, inter_size / pack_size);
+    CHECK_TENSOR_SHAPE(fc3_experts_weights, num_experts, inter_size, hidden_size / pack_size);
+  }
+
+  CHECK_TENSOR_SHAPE(router_probs, num_rows, num_experts);
+
+  CHECK_TENSOR_SHAPE(fc1_experts_bias, num_experts, fc1_inter_size);
+  CHECK_TENSOR_SHAPE(fc2_experts_bias, num_experts, hidden_size);
+  CHECK_TENSOR_SHAPE(fc3_experts_bias, num_experts, inter_size);
+
+  CHECK_TENSOR_SHAPE(fc1_experts_scales, num_experts, fc1_inter_size);
+  CHECK_TENSOR_SHAPE(fc2_experts_scales, num_experts, hidden_size);
+  CHECK_TENSOR_SHAPE(fc3_experts_scales, num_experts, inter_size);
+
+  if (fc3_experts_weights == nullptr) {
+    ORT_ENFORCE(fc3_experts_bias == nullptr && fc3_experts_scales == nullptr);
+  } else { // fc3 exists
+    ORT_ENFORCE(fc1_experts_scales == nullptr || fc3_experts_scales != nullptr); // MOE no scale, or qMOE need scales
+  }
+
+  parameters.num_rows = num_rows;
+  parameters.num_experts = num_experts;
+  parameters.local_num_experts = local_num_experts;
+  parameters.hidden_size = hidden_size;
+  parameters.inter_size = inter_size;
+  if (num_experts == local_num_experts) {
+    if (parameters.tensor_shards == 1) {
+      parameters.parallel_type = MoEParallelType::None;
+    } else {
+      parameters.parallel_type = MoEParallelType::TP;
+    }
+  } else if (num_experts > local_num_experts) {
+    if (parameters.tensor_shards == 1) {
+      parameters.parallel_type = MoEParallelType::EP;
+    } else {
+      parameters.parallel_type = MoEParallelType::EPAndTP;
+    }
+  } else {
+    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                           "num_experts must be greater than or equal to local_num_experts, got ", num_experts,
+                           " and ", local_num_experts);
+  }
+
+  return Status::OK();
+}
+
+}  // namespace moe_helper
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/cuda/collective/sharded_moe.cc

@@ -71,10 +71,14 @@ Status ShardedMoE<T>::ComputeInternal(OpKernelContext* context) const {
   const Tensor* fc3_experts_bias_optional = context->Input<Tensor>(7);
 
   MoEParameters moe_params(tensor_shards_);
-  MoEQuantType quant_type = MoEQuantType::None;
-  ORT_RETURN_IF_ERROR(CheckInputs(moe_params, quant_type, input, router_probs, fc1_experts_weights,
-                                  fc1_experts_bias_optional, fc2_experts_weights, fc2_experts_bias_optional,
-                                  fc3_experts_weights_optional, fc3_experts_bias_optional));
+  MoEParameters moe_params;
+  ORT_RETURN_IF_ERROR(::onnxruntime::contrib::moe_helper::CheckInputs<Tensor>(
+      moe_params, input, router_probs,
+      fc1_experts_weights, fc1_experts_bias_optional, nullptr,
+      fc2_experts_weights, fc2_experts_bias_optional, nullptr,
+      fc3_experts_weights_optional, fc3_experts_bias_optional, nullptr,
+      1,  // no quantization so pack size is 1
+      activation_type_ == ort_fastertransformer::ActivationType::SwiGLU));
 
   ORT_RETURN_IF_NOT(moe_params.num_experts % nccl_->Size() == 0, "num_experts should be divisible by world_size");
 

onnxruntime/contrib_ops/cuda/moe/moe.cc

@@ -39,10 +39,13 @@ Status MoE<T>::ComputeInternal(OpKernelContext* context) const {
   const Tensor* fc3_experts_bias_optional = context->Input<Tensor>(7);
 
   MoEParameters moe_params;
-  MoEQuantType quant_type = MoEQuantType::None;
-  ORT_RETURN_IF_ERROR(CheckInputs(moe_params, quant_type, input, router_probs, fc1_experts_weights,
-                                  fc1_experts_bias_optional, fc2_experts_weights, fc2_experts_bias_optional,
-                                  fc3_experts_weights_optional, fc3_experts_bias_optional));
+  ORT_RETURN_IF_ERROR(::onnxruntime::contrib::moe_helper::CheckInputs<Tensor>(
+      moe_params, input, router_probs,
+      fc1_experts_weights, fc1_experts_bias_optional, nullptr,
+      fc2_experts_weights, fc2_experts_bias_optional, nullptr,
+      fc3_experts_weights_optional, fc3_experts_bias_optional, nullptr,
+      1,  //  no quantization so pack size is 1
+      activation_type_ == ort_fastertransformer::ActivationType::SwiGLU));
 
   using CudaT = typename OrtToCudaType<T>::type;
   auto stream = context->GetComputeStream();