[CPU EP] Add block quantized Gather contrib op (#21630)

### Description
Add a gather that supports block-quantized input data.


### Motivation and Context
To support Web inference scenario with quantized vocabulary embeddings.

docs/ContribOperators.md

@@ -37,6 +37,7 @@ Do not modify directly.*
   * <a href="#com.microsoft.FusedMatMul">com.microsoft.FusedMatMul</a>
   * <a href="#com.microsoft.FusedMatMulActivation">com.microsoft.FusedMatMulActivation</a>
   * <a href="#com.microsoft.GatedRelativePositionBias">com.microsoft.GatedRelativePositionBias</a>
+  * <a href="#com.microsoft.GatherBlockQuantized">com.microsoft.GatherBlockQuantized</a>
   * <a href="#com.microsoft.GatherND">com.microsoft.GatherND</a>
   * <a href="#com.microsoft.Gelu">com.microsoft.Gelu</a>
   * <a href="#com.microsoft.GemmFastGelu">com.microsoft.GemmFastGelu</a>
@@ -2030,6 +2031,64 @@ This version of the operator has been available since version 1 of the 'com.micr
 </dl>
 
 
+### <a name="com.microsoft.GatherBlockQuantized"></a><a name="com.microsoft.gatherblockquantized">**com.microsoft.GatherBlockQuantized**</a>
+
+  GatherBlockQuantized is a Gather with data quantized. It is similar to Gather (https://github.com/onnx/onnx/blob/main/docs/Operators.md#gather) with differences:
+    1. Input `data` is a constant. It is quantized block-wise along attribute `quantize_axis` with block size specified by attribute `block_size`.
+       `block_size must` be a power of 2 and not smaller than 16, like 16, 32, 64, 128, ..
+    2. Input `data`'s scale and zero point are specified by input `scales` and `zero_points`. `scales` and `zero_points` are also constants.
+       If `zero_points` is not provided, 0 is the zero point.
+    3. During the op execution, `data` and `indices` are first used to generate the quantized output. Then, `scales` and `zero_points` are used
+       to dequantize the output.
+    4. The `output` and `scales` have the same type. The `data` and `zero_points` have the same type.
+
+#### Version
+
+This version of the operator has been available since version 1 of the 'com.microsoft' operator set.
+
+#### Attributes
+
+<dl>
+<dt><tt>block_size</tt> : int</dt>
+<dd>(Optional) block size used for weight quantization. It needs to be a power of 2 and not smaller than 16.</dd>
+<dt><tt>gather_axis</tt> : int</dt>
+<dd>(Optional) Which axis to gather on. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).</dd>
+<dt><tt>quantize_axis</tt> : int</dt>
+<dd>(Optional) Which axis to block-wise quantize. Negative value means counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).</dd>
+</dl>
+
+#### Inputs (3 - 4)
+
+<dl>
+<dt><tt>data</tt> : T1</dt>
+<dd>Tensor of rank r >= 1. Block-wise quantized.</dd>
+<dt><tt>indices</tt> : Tind</dt>
+<dd>Tensor of int32/int64 indices, of any rank q. All index values are expected to be within bounds [-s, s-1] along axis of size s. It is an error if any of the index values are out of bounds.</dd>
+<dt><tt>scales</tt> : T2</dt>
+<dd>quantization scale</dd>
+<dt><tt>zero_points</tt> (optional) : T1</dt>
+<dd>quantization zero points</dd>
+</dl>
+
+#### Outputs
+
+<dl>
+<dt><tt>output</tt> : T2</dt>
+<dd>Dequantized output tensor of rank q + (r - 1).</dd>
+</dl>
+
+#### Type Constraints
+
+<dl>
+<dt><tt>T1</tt> : tensor(int4), tensor(uint4)</dt>
+<dd>Constrain quantized types.</dd>
+<dt><tt>T2</tt> : tensor(float), tensor(float16), tensor(bfloat16)</dt>
+<dd>Constrain dequantized types.</dd>
+<dt><tt>Tind</tt> : tensor(int32), tensor(int64)</dt>
+<dd>Constrain indices to integer types.</dd>
+</dl>
+
+
 ### <a name="com.microsoft.GatherND"></a><a name="com.microsoft.gathernd">**com.microsoft.GatherND**</a>
 
   Given `data` tensor of rank r >= 1, and `indices` tensor of rank q >= 1, gather

docs/OperatorKernels.md

@@ -477,6 +477,7 @@ Do not modify directly.*
 |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(float)|
+|GatherBlockQuantized|*in* data:**T1**<br> *in* indices:**Tind**<br> *in* scales:**T2**<br> *in* zero_points:**T1**<br> *out* output:**T2**|1+|**T1** = tensor(int4), tensor(uint4)<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)|
 |Gelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |GreedySearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**I**<br> *in* prefix_vocab_mask:**I**<br> *in* attention_mask:**I**<br> *out* sequences:**I**|1+|**T** = tensor(float)|

onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc

@@ -3,6 +3,7 @@
 
 #include "contrib_ops/cpu/cpu_contrib_kernels.h"
 #include "core/graph/constants.h"
+#include "core/framework/int4.h"
 #include "core/mlas/inc/mlas.h"
 
 namespace onnxruntime {
@@ -33,6 +34,10 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Trans
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FusedMatMul);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulNBits);
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulBnb4);
+class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int32_t, GatherBlockQuantized);
+class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int64_t, GatherBlockQuantized);
+class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Int4x2, int32_t, GatherBlockQuantized);
+class ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Int4x2, int64_t, GatherBlockQuantized);
 #ifndef ORT_MINIMAL_BUILD
 class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulFpQ4);
 #endif
@@ -298,6 +303,10 @@ Status RegisterCpuContribKernels(KernelRegistry& kernel_registry) {
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, FusedMatMul)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulNBits)>,
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulBnb4)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int32_t, GatherBlockQuantized)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, UInt4x2, int64_t, GatherBlockQuantized)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Int4x2, int32_t, GatherBlockQuantized)>,
+      BuildKernelCreateInfo<ONNX_OPERATOR_TWO_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Int4x2, int64_t, GatherBlockQuantized)>,
 #ifndef ORT_MINIMAL_BUILD
       BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, MatMulFpQ4)>,
 #endif

onnxruntime/contrib_ops/cpu/quantization/gather_block_quantized.cc

@@ -0,0 +1,282 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <vector>
+#include <unordered_map>
+
+#include "core/common/common.h"
+#include "core/common/narrow.h"
+#include "core/common/safeint.h"
+#include "core/framework/float16.h"
+#include "core/framework/int4.h"
+#include "core/framework/op_kernel.h"
+#include "core/platform/threadpool.h"
+#include "core/providers/common.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+template <typename T1, typename Tind>
+class GatherBlockQuantized : public OpKernel {
+ public:
+  GatherBlockQuantized(const OpKernelInfo& info) : OpKernel(info) {
+    if (!info.GetAttr<int64_t>("gather_axis", &gather_axis_).IsOK()) {
+      gather_axis_ = 0;
+    }
+
+    if (!info.GetAttr<int64_t>("quantize_axis", &quantize_axis_).IsOK()) {
+      quantize_axis_ = 1;
+    }
+
+    if (!info.GetAttr<int64_t>("block_size", &block_size_).IsOK()) {
+      block_size_ = 128;
+    }
+
+    ORT_ENFORCE(block_size_ >= 16 && ((block_size_ - 1) & block_size_) == 0,
+                "'block_size' must be 2's power and not less than 16.");
+  }
+
+  Status Compute(OpKernelContext* context) const override;
+
+ protected:
+  struct Prepare {
+    const Tensor* data_tensor;
+    const Tensor* indices_tensor;
+    const Tensor* scales_tensor;
+    const Tensor* zero_points_tensor;
+    Tensor* output_tensor;
+    int64_t gather_axis;
+    int64_t quantize_axis;
+  };
+
+  Status PrepareForCompute(OpKernelContext* context, Prepare& args) const;
+
+  template <typename T2>
+  Status CopyDataAndDequantize(const T1* data_ptr,
+                               const Tind* indices_ptr,
+                               const T2* scales_ptr,
+                               const T1* zero_points_ptr,
+                               T2* output_ptr,
+                               const int64_t gather_M,
+                               const int64_t gather_N,
+                               const int64_t gather_axis_dim,
+                               const int64_t gather_block,
+                               const int64_t quantize_axis_dim,
+                               const int64_t quantize_N,
+                               concurrency::ThreadPool* tp) const;
+
+ private:
+  int64_t gather_axis_;
+  int64_t quantize_axis_;
+  int64_t block_size_;
+};
+
+template <typename T1, typename Tind>
+Status GatherBlockQuantized<T1, Tind>::PrepareForCompute(OpKernelContext* context, Prepare& p) const {
+  p.data_tensor = context->Input<Tensor>(0);
+  p.indices_tensor = context->Input<Tensor>(1);
+  p.scales_tensor = context->Input<Tensor>(2);
+  p.zero_points_tensor = context->Input<Tensor>(3);
+
+  const auto& data_shape = p.data_tensor->Shape();
+  const auto& indices_shape = p.indices_tensor->Shape();
+  const auto data_rank = data_shape.NumDimensions();
+  p.gather_axis = HandleNegativeAxis(gather_axis_, narrow<int64_t>(data_rank));
+  p.quantize_axis = HandleNegativeAxis(quantize_axis_, narrow<int64_t>(data_rank));
+
+  std::vector<int64_t> shape;
+  shape.reserve(data_rank - 1 + indices_shape.NumDimensions());
+
+  // get output tensor
+  // replace the dimension for p.gather_axis with the shape from the indices
+  for (int64_t i = 0; i < p.gather_axis; ++i)
+    shape.push_back(data_shape[narrow<size_t>(i)]);
+
+  for (const auto dim : indices_shape.GetDims())
+    shape.push_back(dim);
+
+  for (int64_t i = p.gather_axis + 1; i < static_cast<int64_t>(data_rank); ++i)
+    shape.push_back(data_shape[narrow<size_t>(i)]);
+
+  p.output_tensor = context->Output(0, TensorShape(std::move(shape)));
+
+  // validate quantization parameters
+  const auto& scales_shape = p.scales_tensor->Shape();
+  ORT_RETURN_IF_NOT(data_shape.NumDimensions() == scales_shape.NumDimensions(),
+                    "data and scales must have the same rank.");
+  for (size_t i = 0; i < data_shape.NumDimensions(); ++i) {
+    ORT_RETURN_IF_NOT(i == static_cast<size_t>(p.quantize_axis)
+                          ? (data_shape[i] + block_size_ - 1) / block_size_ == scales_shape[i]
+                          : data_shape[i] == scales_shape[i],
+                      "data and scales do not match shapes.");
+  }
+
+  if (p.zero_points_tensor) {
+    const auto& zero_points_shape = p.zero_points_tensor->Shape();
+    ORT_RETURN_IF_NOT(scales_shape.NumDimensions() == zero_points_shape.NumDimensions(),
+                      "scales and zero_points must have the same rank.");
+    for (size_t i = 0; i < scales_shape.NumDimensions(); ++i) {
+      ORT_RETURN_IF_NOT(scales_shape[i] == zero_points_shape[i],
+                        "scales and zero_points must have the same shape.");
+    }
+  }
+
+  return Status::OK();
+}
+
+template <typename T1, typename Tind>
+template <typename T2>
+Status GatherBlockQuantized<T1, Tind>::CopyDataAndDequantize(const T1* data_ptr,
+                                                             const Tind* indices_ptr,
+                                                             const T2* scales_ptr,
+                                                             const T1* zero_points_ptr,
+                                                             T2* output_ptr,
+                                                             const int64_t gather_M,
+                                                             const int64_t gather_N,
+                                                             const int64_t gather_axis_dim,
+                                                             const int64_t gather_block,
+                                                             const int64_t quantize_axis_dim,
+                                                             const int64_t quantize_N,
+                                                             concurrency::ThreadPool* tp) const {
+  auto data_full_block = gather_axis_dim * gather_block;
+  auto quantize_full_block = quantize_axis_dim * quantize_N;
+  auto scale_full_block = (quantize_axis_dim + block_size_ - 1) / block_size_ * quantize_N;
+
+  auto lambda = [&](int64_t gather_MN_idx, std::unordered_map<int64_t, int64_t>& cache) {
+    int64_t gather_M_idx = gather_MN_idx / gather_N;
+    int64_t gather_N_idx = gather_MN_idx % gather_N;
+
+    int64_t indices_val = static_cast<int64_t>(indices_ptr[gather_N_idx]);
+    ORT_ENFORCE(indices_val >= -gather_axis_dim && indices_val < gather_axis_dim,
+                "indices element out of data bounds, idx=", indices_val,
+                " must be within the inclusive range [", -gather_axis_dim, ",", gather_axis_dim - 1, "]");
+
+    indices_val = indices_val < 0 ? indices_val + gather_axis_dim : indices_val;
+    int64_t output_idx_base = gather_MN_idx * gather_block;
+    int64_t data_idx_base = gather_M_idx * data_full_block + indices_val * gather_block;
+
+    if (auto it = cache.find(data_idx_base); it != cache.end()) {
+      int64_t output_src_idx = it->second;
+      memcpy(output_ptr + output_idx_base, output_ptr + output_src_idx, narrow<size_t>(gather_block * sizeof(T2)));
+      return;
+    }
+
+    // TODO(fajin): use SIMD
+    int64_t output_idx = output_idx_base;
+    int64_t data_idx = data_idx_base;
+    for (int64_t i = 0; i < gather_block; ++i, ++output_idx, ++data_idx) {
+      auto data_val = static_cast<int32_t>(data_ptr[data_idx >> 1].GetElem(narrow<size_t>(data_idx & 1)));
+
+      int64_t x = data_idx / quantize_full_block;
+      int64_t y = data_idx % quantize_full_block / quantize_N;
+      int64_t z = data_idx % quantize_N;
+      int64_t scale_idx = x * scale_full_block + y / block_size_ * quantize_N + z;
+      auto scale_val = static_cast<float>(scales_ptr[scale_idx]);
+      auto zp_val = static_cast<int32_t>(zero_points_ptr
+                                             ? zero_points_ptr[scale_idx >> 1].GetElem(narrow<size_t>(scale_idx & 1))
+                                             : 0);
+
+      output_ptr[output_idx] = static_cast<T2>(static_cast<float>(data_val - zp_val) * scale_val);
+    }
+
+    cache[data_idx_base] = output_idx_base;
+  };
+
+  concurrency::ThreadPool::TryParallelFor(
+      tp,
+      SafeInt<ptrdiff_t>(gather_M) * gather_N,
+      static_cast<double>(gather_block * 3),
+      [&lambda](ptrdiff_t first, ptrdiff_t last) {
+        // cache dequantized gather_block. Key is data_idx_base. Value is the output_idx_base.
+        // cache is per thread to avoid contention.
+        std::unordered_map<int64_t, int64_t> cache;
+
+        for (auto index = static_cast<int64_t>(first), end = static_cast<int64_t>(last);
+             index < end;
+             ++index) {
+          lambda(index, cache);
+        }
+      });
+
+  return Status::OK();
+}
+
+template <typename T1, typename Tind>
+Status GatherBlockQuantized<T1, Tind>::Compute(OpKernelContext* context) const {
+  Prepare p;
+  ORT_RETURN_IF_ERROR(PrepareForCompute(context, p));
+
+  const auto& data_shape = p.data_tensor->Shape();
+  // re-shape the data tensor to [gather_M, gather_axis_dim, gather_block]
+  // re-shape the indices tensor to [gather_N]
+  // re-shape the output tensor to [gather_M, gather_N, gather_block]
+  // For an index i in the output tensor:
+  //  1> the output block index is blk_i = i / gather_block, block element index is blk_ele_i = i % gather_block,
+  //  2> block is picked from data based on value from indices: axis_i = indices[blk_i % gather_N],
+  //  3> get the corresponding block in data tensor: data_blk = data[blk_i / gather_N, axis_i, :],
+  //  4> pick the element from the block: value_i = data_blk[blk_ele_i]
+  const int64_t gather_block = data_shape.SizeFromDimension(SafeInt<size_t>(p.gather_axis) + 1);
+  const int64_t gather_axis_dim = data_shape[narrow<size_t>(p.gather_axis)];
+  const int64_t gather_M = data_shape.SizeToDimension(narrow<size_t>(p.gather_axis));
+  const int64_t gather_N = p.indices_tensor->Shape().Size();
+  // re-shape the data tensor to [quantize_M, quantize_axis_dim, quantize_N]
+  // For an index i in the output tensor:
+  //  1> based on previous comment, corresponding data index is (blk_i / gather_N, axis_i, blk_ele_i)
+  //  2> flatten the data index:
+  //     data_i = blk_i / gather_N * gather_axis_dim * gather_block + axis_i * gather_block + blk_ele_i
+  //  3> map data_i to quantize shape: (x, y, z) =
+  //     (data_i / (quantize_axis_dim * quantize_N),
+  //      data_i % (quantize_axis_dim * quantize_N) / quantize_N,
+  //      data_i % quantize_N)
+  //  4> get scale index: (x, y / block_size_, z)
+  const int64_t quantize_axis_dim = data_shape[narrow<size_t>(p.quantize_axis)];
+  const int64_t quantize_N = data_shape.SizeFromDimension(SafeInt<size_t>(p.quantize_axis) + 1);
+
+  concurrency::ThreadPool* tp = context->GetOperatorThreadPool();
+  const auto* data_ptr = p.data_tensor->template Data<T1>();
+  const auto* indices_ptr = p.indices_tensor->template Data<Tind>();
+  const auto* zero_points_ptr = p.zero_points_tensor ? p.zero_points_tensor->template Data<T1>() : nullptr;
+  const auto dequantized_type = p.scales_tensor->GetElementType();
+
+  if (dequantized_type == ONNX_NAMESPACE::TensorProto::FLOAT) {
+    const auto* scales_ptr = p.scales_tensor->template Data<float>();
+    auto* output_ptr = p.output_tensor->template MutableData<float>();
+
+    return CopyDataAndDequantize<float>(data_ptr, indices_ptr, scales_ptr, zero_points_ptr,
+                                        output_ptr, gather_M, gather_N, gather_axis_dim, gather_block,
+                                        quantize_axis_dim, quantize_N,
+                                        tp);
+  } else if (dequantized_type == ONNX_NAMESPACE::TensorProto::FLOAT16) {
+    const auto* scales_ptr = p.scales_tensor->template Data<MLFloat16>();
+    auto* output_ptr = p.output_tensor->template MutableData<MLFloat16>();
+
+    return CopyDataAndDequantize<MLFloat16>(data_ptr, indices_ptr, scales_ptr, zero_points_ptr,
+                                            output_ptr, gather_M, gather_N, gather_axis_dim, gather_block,
+                                            quantize_axis_dim, quantize_N,
+                                            tp);
+  } else if (dequantized_type == ONNX_NAMESPACE::TensorProto::BFLOAT16) {
+    ORT_THROW("DequantizeLinear into BFLOAT16 is not implemented yet.");
+  } else {
+    ORT_THROW("Unsupported dequantized type: ", dequantized_type);
+  }
+}
+
+#define REGISTER_GATHERBLOCKQUANTIZED(T1, Tind)                                                                   \
+  ONNX_OPERATOR_TWO_TYPED_KERNEL_EX(                                                                              \
+      GatherBlockQuantized,                                                                                       \
+      kMSDomain, 1,                                                                                               \
+      T1, Tind,                                                                                                   \
+      kCpuExecutionProvider,                                                                                      \
+      KernelDefBuilder()                                                                                          \
+          .TypeConstraint("T1", DataTypeImpl::GetTensorType<T1>())                                                \
+          .TypeConstraint("T2", {DataTypeImpl::GetTensorType<float>(), DataTypeImpl::GetTensorType<MLFloat16>()}) \
+          .TypeConstraint("Tind", DataTypeImpl::GetTensorType<Tind>()),                                           \
+      GatherBlockQuantized<T1, Tind>);
+
+REGISTER_GATHERBLOCKQUANTIZED(UInt4x2, int32_t);
+REGISTER_GATHERBLOCKQUANTIZED(UInt4x2, int64_t);
+REGISTER_GATHERBLOCKQUANTIZED(Int4x2, int32_t);
+REGISTER_GATHERBLOCKQUANTIZED(Int4x2, int64_t);
+
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -3544,6 +3544,124 @@ MatMulBnb4 is a MatMul with weight quantized with 4 bits using either FP4 or NF4
         MatmulWithQuantWeightShapeInference(ctx, in_features, out_features, transB);
       });
 
+  static const char* GatherBlockQuantized_ver1_doc = R"DOC(
+GatherBlockQuantized is a Gather with data quantized. It is similar to Gather (https://github.com/onnx/onnx/blob/main/docs/Operators.md#gather) with differences:
+  1. Input `data` is a constant. It is quantized block-wise along attribute `quantize_axis` with block size specified by attribute `block_size`.
+     `block_size must` be a power of 2 and not smaller than 16, like 16, 32, 64, 128, ..
+  2. Input `data`'s scale and zero point are specified by input `scales` and `zero_points`. `scales` and `zero_points` are also constants.
+     If `zero_points` is not provided, 0 is the zero point.
+  3. During the op execution, `data` and `indices` are first used to generate the quantized output. Then, `scales` and `zero_points` are used
+     to dequantize the output.
+  4. The `output` and `scales` have the same type. The `data` and `zero_points` have the same type.
+)DOC";
+
+  ONNX_CONTRIB_OPERATOR_SCHEMA(GatherBlockQuantized)
+      .SetDomain(kMSDomain)
+      .SinceVersion(1)
+      .SetDoc(GatherBlockQuantized_ver1_doc)
+      .Attr("gather_axis",
+            "(Optional) Which axis to gather on. Negative value means "
+            "counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).",
+            AttributeProto::INT, static_cast<int64_t>(0))
+      .Attr("quantize_axis",
+            "(Optional) Which axis to block-wise quantize. Negative value means "
+            "counting dimensions from the back. Accepted range is [-r, r-1] where r = rank(data).",
+            AttributeProto::INT, static_cast<int64_t>(1))
+      .Attr("block_size",
+            "(Optional) block size used for weight quantization. It needs to be a power of 2 and not smaller than 16.",
+            AttributeProto::INT,
+            static_cast<int64_t>(128))
+      .Input(0, "data", "Tensor of rank r >= 1. Block-wise quantized.", "T1")
+      .Input(1,
+             "indices",
+             "Tensor of int32/int64 indices, of any rank q. All index values are expected to be within bounds [-s, s-1] "
+             "along axis of size s. It is an error if any of the index values are out of bounds.",
+             "Tind")
+      .Input(2, "scales", "quantization scale", "T2")
+      .Input(3, "zero_points", "quantization zero points", "T1", OpSchema::Optional)
+      .Output(0, "output", "Dequantized output tensor of rank q + (r - 1).", "T2")
+      .TypeConstraint("T1", {"tensor(int4)", "tensor(uint4)"}, "Constrain quantized types.")
+      .TypeConstraint("T2", {"tensor(float)", "tensor(float16)", "tensor(bfloat16)"}, "Constrain dequantized types.")
+      .TypeConstraint("Tind", {"tensor(int32)", "tensor(int64)"}, "Constrain indices to integer types.")
+      .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
+        // Type inference
+        propagateElemTypeFromInputToOutput(ctx, 2, 0);
+        // Shape inference
+        if (!hasNInputShapes(ctx, 3)) {
+          return;
+        }
+        const TensorShapeProto& data_shape = ctx.getInputType(0)->tensor_type().shape();
+        const TensorShapeProto& indices_shape = ctx.getInputType(1)->tensor_type().shape();
+        const TensorShapeProto& scales_shape = ctx.getInputType(2)->tensor_type().shape();
+        int r = data_shape.dim_size();
+
+        if (r < 1) {
+          fail_shape_inference("data tensor must have rank >= 1");
+        }
+
+        int gather_axis = static_cast<int>(getAttribute(ctx, "gather_axis", 0));
+        int quantize_axis = static_cast<int>(getAttribute(ctx, "quantize_axis", 1));
+        auto block_size = getAttribute(ctx, "block_size", 128);
+        if (gather_axis < -r || gather_axis >= r) {
+          fail_shape_inference("gather_axis must be in [-r, r-1]");
+        }
+        if (quantize_axis < -r || quantize_axis >= r) {
+          fail_shape_inference("quantize_axis must be in [-r, r-1]");
+        }
+        if (block_size < 0) {
+          fail_shape_inference("block_size must be non-negative");
+        }
+
+        gather_axis = (gather_axis + r) % r;
+        quantize_axis = (quantize_axis + r) % r;
+
+        if (scales_shape.dim_size() != r) {
+          fail_shape_inference("scales must have the same rank as data");
+        }
+
+        for (int i = 0; i < r; ++i) {
+          if (!data_shape.dim(i).has_dim_value() ||
+              !scales_shape.dim(i).has_dim_value() ||
+              (i == quantize_axis && (data_shape.dim(i).dim_value() + block_size - 1) / block_size != scales_shape.dim(i).dim_value()) ||
+              (i != quantize_axis && data_shape.dim(i).dim_value() != scales_shape.dim(i).dim_value())) {
+            fail_shape_inference("data shape and scales shape do not match");
+          }
+        }
+
+        // validate zero point shape
+        if (ctx.hasInput(3)) {
+          if (!hasInputShape(ctx, 3)) {
+            fail_shape_inference("zero_points shape must be known");
+          }
+
+          const auto& zp_shape = getInputShape(ctx, 3);
+          if (zp_shape.dim_size() != r) {
+            fail_shape_inference("zero points must have the same rank as data");
+          }
+
+          for (int i = 0; i < r; ++i) {
+            if (!zp_shape.dim(i).has_dim_value() ||
+                zp_shape.dim(i).dim_value() != scales_shape.dim(i).dim_value()) {
+              fail_shape_inference("zero points shape and scales shape do not match");
+            }
+          }
+        }
+
+        int q = indices_shape.dim_size();
+        int out_rank = q + r - 1;
+        if (out_rank == 0) {
+          ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape();
+        }
+        for (int i = 0; i < out_rank; ++i) {
+          *ctx.getOutputType(0)->mutable_tensor_type()->mutable_shape()->add_dim() =
+              (i < gather_axis)
+                  ? data_shape.dim(i)
+              : (i >= gather_axis && i < gather_axis + q)
+                  ? indices_shape.dim(i - gather_axis)
+                  : data_shape.dim(i - q + 1);
+        }
+      });
+
 #ifdef ENABLE_ATEN
   ONNX_CONTRIB_OPERATOR_SCHEMA(ATen)
       .SetDomain(kPytorchAtenDomain)

onnxruntime/test/contrib_ops/gather_block_quantized_op_test.cc

@@ -0,0 +1,468 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include <vector>
+#include <type_traits>
+#include <memory>
+#include <utility>
+
+#include "core/common/common.h"
+#include "core/framework/execution_provider.h"
+#include "gtest/gtest.h"
+#include "test/providers/provider_test_utils.h"
+#include "test/util/include/default_providers.h"
+
+namespace onnxruntime {
+namespace test {
+
+// Combinations: types, gather_axis, quantize_axis, block_size, indices, scale shape vs data shape
+template <typename T1, typename T2, typename Tind>
+void RunGatherBlockQuantized(const std::vector<T1>& data,
+                             const std::vector<int64_t>& data_shape,
+                             const std::vector<Tind>& indices,
+                             const std::vector<int64_t>& indices_shape,
+                             const std::vector<T2>& scales,
+                             const std::vector<int64_t>& scales_shape,
+                             const std::vector<T1>& zero_points,
+                             const int64_t gather_axis,
+                             const int64_t quantize_axis,
+                             const int64_t block_size,
+                             const std::vector<T2>& output,
+                             const std::vector<int64_t>& output_shape,
+                             OpTester::ExpectResult expect_result = OpTester::ExpectResult::kExpectSuccess) {
+  auto run_test = [&](bool indices_is_initializer) {
+    OpTester test("GatherBlockQuantized", 1, kMSDomain);
+
+    test.AddAttribute<int64_t>("gather_axis", gather_axis);
+    test.AddAttribute<int64_t>("quantize_axis", quantize_axis);
+    test.AddAttribute<int64_t>("block_size", block_size);
+
+    test.AddInput<T1>("data", data_shape, data);
+    test.AddInput<Tind>("indices", indices_shape, indices, indices_is_initializer);
+    test.AddInput<T2>("scales", scales_shape, scales);
+    if (!zero_points.empty()) {
+      test.AddInput<T1>("zero_points", scales_shape, zero_points);
+    }
+
+    test.AddOutput<T2>("output", output_shape, output);
+
+    std::vector<std::unique_ptr<IExecutionProvider>> eps;
+    eps.push_back(DefaultCpuExecutionProvider());
+    test.Run(expect_result, "", {}, nullptr, &eps);
+  };
+
+  run_test(false);
+  run_test(true);
+}
+
+template <typename T1, typename T2>
+typename std::enable_if<
+    (boost::mp11::mp_contains<TypeList<BFloat16, MLFloat16, float>, T1>::value && std::is_same<T2, float>::value) ||
+        (std::is_integral<T1>::value && std::is_same<T2, int>::value),
+    std::vector<T1>>::type
+ToType(const std::vector<T2>& vec) {
+  std::vector<T1> result;
+  for (auto v : vec) {
+    result.push_back(static_cast<T1>(v));
+  }
+
+  return result;
+}
+
+template <typename T>
+typename std::enable_if<boost::mp11::mp_contains<TypeList<UInt4x2, Int4x2>, T>::value, std::vector<T>>::type
+ToType(const std::vector<int>& vec) {
+  std::vector<T> result;
+  size_t i = 0;
+  constexpr int offset = std::is_same<T, Int4x2>::value ? 0 : 8;
+  for (i = 0; i + 1 < vec.size(); i += 2) {
+    result.push_back(T(vec[i] + offset, vec[i + 1] + offset));
+  }
+  if (i < vec.size()) {
+    result.push_back(T(vec[i] + offset, 0 + offset));
+  }
+
+  return result;
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_Fail_WithZeroPoints(int64_t gather_axis,
+                              int64_t quantize_axis,
+                              int64_t block_size) {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {1};
+  std::vector<int64_t> indices_shape = {1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {2, 3, 1};
+  std::vector<int> zero_points = {-1, 1, 0, 0, 1, -1};
+  std::vector<float> output = {8.f, 10.f, 12.f, 14.f,
+                               3.f, 4.f, 5.f, 6.f,
+                               -6.f, -4.f, -2.f, 0.f};
+  std::vector<int64_t> output_shape = {1, 3, 4};
+
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          gather_axis,
+                          quantize_axis,
+                          block_size,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectFailure);
+}
+
+TEST(GatherBlockQuantizedOpTest, UnsupportedTypes) {
+  Test_Fail_WithZeroPoints<int8_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<uint8_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<int16_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<uint16_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<int32_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<uint32_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<int64_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<uint64_t, float, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<UInt4x2, float, int16_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<Int4x2, float, int16_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<UInt4x2, BFloat16, int32_t>(0, 2, 16);
+  Test_Fail_WithZeroPoints<Int4x2, BFloat16, int32_t>(0, 2, 16);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidBlockSize) {
+  Test_Fail_WithZeroPoints<UInt4x2, float, int32_t>(0, 2, 8);
+  Test_Fail_WithZeroPoints<Int4x2, float, int32_t>(0, 2, 17);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidGatherAxis) {
+  Test_Fail_WithZeroPoints<UInt4x2, float, int32_t>(3, 2, 16);
+  Test_Fail_WithZeroPoints<Int4x2, float, int32_t>(-4, 2, 16);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidQuantizeAxis) {
+  Test_Fail_WithZeroPoints<UInt4x2, float, int32_t>(0, 3, 16);
+  Test_Fail_WithZeroPoints<Int4x2, float, int32_t>(0, -4, 16);
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_ShapeMismatch_WithZeroPoints() {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {1};
+  std::vector<int64_t> indices_shape = {1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {2, 2, 1};
+  std::vector<int> zero_points = {-1, 1, 0, 0};
+  std::vector<float> output = {8.f, 10.f, 12.f, 14.f,
+                               3.f, 4.f, 5.f, 6.f,
+                               -6.f, -4.f, -2.f, 0.f};
+  std::vector<int64_t> output_shape = {1, 3, 4};
+
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          0,
+                          2,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectFailure);
+}
+
+TEST(GatherBlockQuantizedOpTest, ShapeMismatch) {
+  Test_ShapeMismatch_WithZeroPoints<UInt4x2, float, int32_t>();
+  Test_ShapeMismatch_WithZeroPoints<Int4x2, float, int32_t>();
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_InvalidIndices_WithZeroPoints() {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {2};
+  std::vector<int64_t> indices_shape = {1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {2, 3, 1};
+  std::vector<int> zero_points = {-1, 1, 0, 0, 1, -1};
+  std::vector<float> output = {8.f, 10.f, 12.f, 14.f,
+                               3.f, 4.f, 5.f, 6.f,
+                               -6.f, -4.f, -2.f, 0.f};
+  std::vector<int64_t> output_shape = {1, 3, 4};
+
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          0,
+                          2,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectFailure);
+}
+
+TEST(GatherBlockQuantizedOpTest, InvalidIndices) {
+  Test_InvalidIndices_WithZeroPoints<UInt4x2, float, int32_t>();
+  Test_InvalidIndices_WithZeroPoints<Int4x2, float, int32_t>();
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_GatherAxis0_WithZeroPoints() {
+  std::vector<int> data = {-8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8, -7, -6, -5, -8,
+                           -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4,
+                           0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0,
+                           4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4,
+                           4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4, 5, 6, 7, 4,
+                           -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4, -3, -2, -1, -4};
+  std::vector<int64_t> data_shape = {2, 3, 17};
+  std::vector<int> indices = {1};
+  std::vector<int64_t> indices_shape = {1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f,
+                               2.0f, 2.0f, 1.0f, 1.0f, 2.0f, 1.0f};
+  std::vector<int64_t> scales_shape = {2, 3, 2};
+  std::vector<int> zero_points = {-1, 1, 0, 0, 1, -1,
+                                  1, -1, 1, 0, -1, 1};
+  std::vector<float> output = {6, 8, 10, 12, 6, 8, 10, 12, 6, 8, 10, 12, 6, 8, 10, 12, 10,
+                               3, 4, 5, 6, 3, 4, 5, 6, 3, 4, 5, 6, 3, 4, 5, 6, 4,
+                               -6, -4, -2, 0, -6, -4, -2, 0, -6, -4, -2, 0, -6, -4, -2, 0, -5};
+  std::vector<int64_t> output_shape = {1, 3, 17};
+
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          0,
+                          2,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          -3,
+                          -1,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+}
+
+TEST(GatherBlockQuantizedOpTest, GatherAxis0WithZeroPoints) {
+  Test_GatherAxis0_WithZeroPoints<UInt4x2, float, int32_t>();
+  Test_GatherAxis0_WithZeroPoints<Int4x2, float, int32_t>();
+  Test_GatherAxis0_WithZeroPoints<UInt4x2, MLFloat16, int32_t>();
+  Test_GatherAxis0_WithZeroPoints<Int4x2, MLFloat16, int32_t>();
+  Test_GatherAxis0_WithZeroPoints<UInt4x2, float, int64_t>();
+  Test_GatherAxis0_WithZeroPoints<Int4x2, float, int64_t>();
+  Test_GatherAxis0_WithZeroPoints<UInt4x2, MLFloat16, int64_t>();
+  Test_GatherAxis0_WithZeroPoints<Int4x2, MLFloat16, int64_t>();
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_GatherAxis0_NoZeroPoints() {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {1};
+  std::vector<int64_t> indices_shape = {1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {2, 3, 1};
+  std::vector<float> output = {8.f, 10.f, 12.f, 14.f,
+                               4.f, 5.f, 6.f, 7.f,
+                               -8.f, -6.f, -4.f, -2.f};
+  std::vector<int64_t> output_shape = {1, 3, 4};
+
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          {},
+                          0,
+                          2,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          {},
+                          -3,
+                          -1,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+}
+
+TEST(GatherBlockQuantizedOpTest, GatherAxis0NoZeroPoints) {
+  Test_GatherAxis0_NoZeroPoints<Int4x2, float, int32_t>();
+  Test_GatherAxis0_NoZeroPoints<Int4x2, MLFloat16, int32_t>();
+  Test_GatherAxis0_NoZeroPoints<Int4x2, float, int64_t>();
+  Test_GatherAxis0_NoZeroPoints<Int4x2, MLFloat16, int64_t>();
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_GatherAxis1_WithZeroPoints() {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {2, -3, 2};
+  std::vector<int64_t> indices_shape = {1, 3};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {2, 1, 4};
+  std::vector<int> zero_points = {-1, 1, 0, 0, 1, -1, 0, 0};
+  std::vector<float> output = {1.f, 0.f, 2.f, 6.f,
+                               -7.f, -16.f, -6.f, -10.f,
+                               1.f, 0.f, 2.f, 6.f,
+                               -5.f, -4.f, -2.f, -2.f,
+                               3.f, 12.f, 6.f, 14.f,
+                               -5.f, -4.f, -2.f, -2.f};
+  std::vector<int64_t> output_shape = {2, 1, 3, 4};
+
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          1,
+                          1,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          -2,
+                          -2,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+}
+
+TEST(GatherBlockQuantizedOpTest, GatherAxis1) {
+  Test_GatherAxis1_WithZeroPoints<UInt4x2, float, int32_t>();
+  Test_GatherAxis1_WithZeroPoints<Int4x2, float, int32_t>();
+  Test_GatherAxis1_WithZeroPoints<UInt4x2, MLFloat16, int32_t>();
+  Test_GatherAxis1_WithZeroPoints<Int4x2, MLFloat16, int32_t>();
+  Test_GatherAxis1_WithZeroPoints<UInt4x2, float, int64_t>();
+  Test_GatherAxis1_WithZeroPoints<Int4x2, float, int64_t>();
+  Test_GatherAxis1_WithZeroPoints<UInt4x2, MLFloat16, int64_t>();
+  Test_GatherAxis1_WithZeroPoints<Int4x2, MLFloat16, int64_t>();
+}
+
+template <typename T1, typename T2, typename Tind>
+void Test_GatherAxis2_WithZeroPoints() {
+  std::vector<int> data = {-8, -7, -6, -5,
+                           -4, -3, -2, -1,
+                           0, 1, 2, 3,
+                           4, 5, 6, 7,
+                           4, 5, 6, 7,
+                           -4, -3, -2, -1};
+  std::vector<int64_t> data_shape = {2, 3, 4};
+  std::vector<int> indices = {-2, 0};
+  std::vector<int64_t> indices_shape = {2, 1};
+  std::vector<float> scales = {1.0f, 2.0f, 1.0f, 2.0f,
+                               1.0f, 2.0f, 1.0f, 2.0f,
+                               1.0f, 2.0f, 1.0f, 2.0f};
+  std::vector<int64_t> scales_shape = {1, 3, 4};
+  std::vector<int> zero_points = {-1, 1, 0, 0,
+                                  1, -1, 0, 0,
+                                  0, 0, 1, -1};
+  std::vector<float> output = {-6.f, -7.f, -2.f, -5.f, 1.f, 0.f,
+                               6.f, 5.f, 6.f, 3.f, -3.f, -4.f};
+  std::vector<int64_t> output_shape = {2, 3, 2, 1};
+
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          2,
+                          0,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+  RunGatherBlockQuantized(ToType<T1>(data),
+                          data_shape,
+                          ToType<Tind>(indices),
+                          indices_shape,
+                          ToType<T2>(scales),
+                          scales_shape,
+                          ToType<T1>(zero_points),
+                          -1,
+                          -3,
+                          16,
+                          ToType<T2>(output),
+                          output_shape,
+                          OpTester::ExpectResult::kExpectSuccess);
+}
+
+TEST(GatherBlockQuantizedOpTest, GatherAxis2) {
+  Test_GatherAxis2_WithZeroPoints<UInt4x2, float, int32_t>();
+  Test_GatherAxis2_WithZeroPoints<Int4x2, float, int32_t>();
+  Test_GatherAxis2_WithZeroPoints<UInt4x2, MLFloat16, int32_t>();
+  Test_GatherAxis2_WithZeroPoints<Int4x2, MLFloat16, int32_t>();
+  Test_GatherAxis2_WithZeroPoints<UInt4x2, float, int64_t>();
+  Test_GatherAxis2_WithZeroPoints<Int4x2, float, int64_t>();
+  Test_GatherAxis2_WithZeroPoints<UInt4x2, MLFloat16, int64_t>();
+  Test_GatherAxis2_WithZeroPoints<Int4x2, MLFloat16, int64_t>();
+}
+
+}  // namespace test
+}  // namespace onnxruntime