diff --git a/src/common/cuda_utils.h b/src/common/cuda_utils.h
index 2f7254040475..ccf0931f2480 100644
--- a/src/common/cuda_utils.h
+++ b/src/common/cuda_utils.h
@@ -187,6 +187,69 @@ namespace mxnet {
 namespace common {
 /*! \brief common utils for cuda */
 namespace cuda {
+/*!
+ * \brief Converts between C++ datatypes and enums/constants needed by cuBLAS.
+ */
+template<typename DType>
+struct CublasType;
+
+// With CUDA v8, cuBLAS adopted use of cudaDataType_t instead of its own
+// datatype cublasDataType_t.  The older cudaDataType_t values could be
+// included below, but since this class was introduced to support the cuBLAS v8
+// call cublasGemmEx(), burdening the class with the legacy type values
+// was not needed.
+
+template<>
+struct CublasType<float> {
+  static const int kFlag = mshadow::kFloat32;
+#if CUDA_VERSION >= 8000
+  static const cudaDataType_t kCudaFlag = CUDA_R_32F;
+#endif
+  typedef float ScaleType;
+  static const float one;
+  static const float zero;
+};
+template<>
+struct CublasType<double> {
+  static const int kFlag = mshadow::kFloat64;
+#if CUDA_VERSION >= 8000
+  static const cudaDataType_t kCudaFlag = CUDA_R_64F;
+#endif
+  typedef double ScaleType;
+  static const double one;
+  static const double zero;
+};
+template<>
+struct CublasType<mshadow::half::half_t> {
+  static const int kFlag = mshadow::kFloat16;
+#if CUDA_VERSION >= 8000
+  static const cudaDataType_t kCudaFlag = CUDA_R_16F;
+#endif
+  typedef float ScaleType;
+  static const mshadow::half::half_t one;
+  static const mshadow::half::half_t zero;
+};
+template<>
+struct CublasType<uint8_t> {
+  static const int kFlag = mshadow::kUint8;
+#if CUDA_VERSION >= 8000
+  static const cudaDataType_t kCudaFlag = CUDA_R_8I;
+#endif
+  typedef uint8_t ScaleType;
+  static const uint8_t one = 1;
+  static const uint8_t zero = 0;
+};
+template<>
+struct CublasType<int32_t> {
+  static const int kFlag = mshadow::kInt32;
+#if CUDA_VERSION >= 8000
+  static const cudaDataType_t kCudaFlag = CUDA_R_32I;
+#endif
+  typedef int32_t ScaleType;
+  static const int32_t one = 1;
+  static const int32_t zero = 0;
+};
+
 /*!
  * \brief Get string representation of cuBLAS errors.
  * \param error The error.
@@ -218,6 +281,17 @@ inline const char* CublasGetErrorString(cublasStatus_t error) {
   return "Unknown cuBLAS status";
 }
 
+#if CUDA_VERSION >= 8000
+/*!
+ * \brief Create the proper constant for indicating cuBLAS transposition, if desired.
+ * \param transpose Whether transposition should be performed.
+ * \return the yes/no transposition-indicating constant.
+ */
+inline cublasOperation_t CublasTransposeOp(bool transpose) {
+  return transpose ? CUBLAS_OP_T : CUBLAS_OP_N;
+}
+#endif
+
 /*!
  * \brief Get string representation of cuSOLVER errors.
  * \param error The error.
diff --git a/src/operator/contrib/transformer-inl.h b/src/operator/contrib/transformer-inl.h
index da3d14e33cf4..da48ffa52dca 100644
--- a/src/operator/contrib/transformer-inl.h
+++ b/src/operator/contrib/transformer-inl.h
@@ -34,6 +34,15 @@
 namespace mxnet {
 namespace op {
 
+struct InterleavedMatMulParam : public dmlc::Parameter<InterleavedMatMulParam> {
+  int heads;
+  bool bwd_ignore_zero_init;
+  DMLC_DECLARE_PARAMETER(InterleavedMatMulParam) {
+    DMLC_DECLARE_FIELD(heads)
+    .describe("Set number of heads");
+  }
+};
+
 template<typename xpu>
 static void DivSqrtDimForward_(const nnvm::NodeAttrs& attrs,
                   const OpContext& ctx,
diff --git a/src/operator/contrib/transformer.cc b/src/operator/contrib/transformer.cc
index 00085c0dc7aa..2ca6f8c71093 100644
--- a/src/operator/contrib/transformer.cc
+++ b/src/operator/contrib/transformer.cc
@@ -29,6 +29,276 @@
 namespace mxnet {
 namespace op {
 
+DMLC_REGISTER_PARAMETER(InterleavedMatMulParam);
+
+static bool InterleavedMatMulSelfAttQKShape(const NodeAttrs& attrs,
+                                            mxnet::ShapeVector* in_shape,
+                                            mxnet::ShapeVector* out_shape) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  CHECK_EQ(in_shape->size(), 1U) << "Input:[queries_keys_values] currently have, "
+                                 << in_shape->size() << " inputs";
+  auto qkv_shape = in_shape->at(0);
+  CHECK_EQ(qkv_shape.ndim(), 3U)
+    << "Input queries_keys_values should be 3D in seq_length-batch-proj_dim, "
+    << "currently is: " << qkv_shape.ndim() << "D";
+  out_shape->resize(1);
+  SHAPE_ASSIGN_CHECK(*out_shape, 0,
+    mxnet::TShape({params.heads * qkv_shape[1], qkv_shape[0], qkv_shape[0]}));
+  return true;
+}
+
+static bool InterleavedMatMulSelfAttValAttShape(const NodeAttrs& attrs,
+                                                mxnet::ShapeVector* in_shape,
+                                                mxnet::ShapeVector* out_shape) {
+  CHECK_EQ(in_shape->size(), 2U) << "Input:[queries_keys_values, attention] currently have, "
+                                 << in_shape->size() << " inputs";
+  auto qkv_shape = in_shape->at(0);
+  auto att_shape = in_shape->at(1);
+  CHECK_EQ(qkv_shape.ndim(), 3U)
+    << "Input queries_keys_values should be 3D in seq_length-batch-3*proj_dim, "
+    << "currently is: " << qkv_shape.ndim() << "D";
+  CHECK_EQ(att_shape.ndim(), 3U)
+    << "Input attention should be 3D in batch-seq_length-seq_length, "
+    << "currently is: " << att_shape.ndim() << "D";
+  CHECK_EQ(qkv_shape[0], att_shape[1])
+    << "queries_keys_values.shape[0] and attention.shape[1] should be the same, "
+    << "currently are " << qkv_shape[0] << " and " << att_shape[1];
+  CHECK_EQ(qkv_shape[0], att_shape[2])
+    << "queries_keys_values.shape[0] and attention.shape[2] should be the same, "
+    << "currently are " << qkv_shape[0] << " and " << att_shape[2];
+  CHECK_EQ(qkv_shape[2] % 3, 0)
+    << "queries_keys_values.shape[2] should be a multiple of 3, "
+    << "currently is " << qkv_shape[2];
+  SHAPE_ASSIGN_CHECK(*out_shape, 0,
+    mxnet::TShape({qkv_shape[0], qkv_shape[1], qkv_shape[2] / 3}));
+  return true;
+}
+
+static bool InterleavedMatMulEncDecQKShape(const NodeAttrs& attrs,
+                                           mxnet::ShapeVector* in_shape,
+                                           mxnet::ShapeVector* out_shape) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  CHECK_EQ(in_shape->size(), 2U) << "Input:[queries, keys_values], currently have "
+                                 << in_shape->size() << " inputs";
+  auto q_shape = in_shape->at(0);
+  auto kv_shape = in_shape->at(1);
+  CHECK_EQ(q_shape.ndim(), 3U) << "Input queries should be 3D in seq_length-batch-proj_dim, "
+                               << "currently is " << q_shape.ndim() << "D";
+  CHECK_EQ(kv_shape.ndim(), 3U) << "Input queries should be 3D in seq_length-batch-2*proj_dim, "
+                                << "currently is " << kv_shape.ndim() << "D";
+  CHECK_EQ(q_shape[2] * 2, kv_shape[2])
+    << "keys_values.shape[2] should be equal to queries.shape[2] * 2, "
+    << "currently are: " << kv_shape[2] << " and " << q_shape[2];
+  CHECK_EQ(q_shape[1], kv_shape[1])
+    << "queries.shape[1] should be equal to keys_values.shape[1], "
+    << "currently are: " << q_shape[1] << " and " << kv_shape[1];
+  SHAPE_ASSIGN_CHECK(*out_shape, 0,
+      mxnet::TShape({q_shape[1] * params.heads, q_shape[0], kv_shape[0]}));
+  return true;
+}
+
+static bool InterleavedMatMulEncDecValAttShape(const NodeAttrs& attrs,
+                                               mxnet::ShapeVector* in_shape,
+                                               mxnet::ShapeVector* out_shape) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  CHECK_EQ(in_shape->size(), 2U) << "Input: [keys_values, attention], currently have "
+                                 << in_shape->size() << " inputs";
+  auto kv_shape = in_shape->at(0);
+  auto att_shape = in_shape->at(1);
+  CHECK_EQ(kv_shape.ndim(), 3U)
+    << "Input keys_values should be 3D in seq_length-batch-2*proj_dim, "
+    << "currently is " << kv_shape.ndim() << "D";
+  CHECK_EQ(att_shape.ndim(), 3U)
+    << "Input attention should be 3D in batch-seq_length-seq_length, "
+    << "currently is " << att_shape.ndim() << "D";
+  CHECK_EQ(kv_shape[0], att_shape[2])
+    << "keys_values.shape[0] should be equal to attention.shape[2], currently are "
+    << kv_shape[0] << " and " << att_shape[2];
+  CHECK_EQ(kv_shape[1] * params.heads, att_shape[0]) << "attention.shape[0] "
+    << "should be equal to keys_values.shape[1] * heads, currently are: "
+    << att_shape[2] << " and " << kv_shape[1];
+  SHAPE_ASSIGN_CHECK(*out_shape, 0,
+      mxnet::TShape({att_shape[1], kv_shape[1], kv_shape[2] / 2}));
+  return true;
+}
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_selfatt_qk)
+.describe(R"code(Compute the matrix multiplication between the projections of
+queries and keys in multihead attention use as self attention.
+
+the input must be a single tensor of interleaved projections
+of queries, keys and values following the layout:
+(seq_length, batch_size, num_heads * head_dim * 3)
+
+the equivalent code would be:
+tmp = mx.nd.reshape(queries_keys_values, shape=(0, 0, num_heads, 3, -1))
+q_proj = mx.nd.transpose(tmp[:,:,:,0,:], axes=(1, 2, 0, 3))
+q_proj = mx.nd.reshape(q_proj, shape=(-1, 0, 0), reverse=True)
+q_proj = mx.nd.contrib.div_sqrt_dim(q_proj)
+k_proj = mx.nd.transpose(tmp[:,:,:,1,:], axes=(1, 2, 0, 3))
+k_proj = mx.nd.reshap(k_proj, shape=(-1, 0, 0), reverse=True)
+output = mx.nd.batch_dot(q_proj, k_proj, transpose_b=True)
+
+This Op is GPU only
+)code" ADD_FILELINE)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>)
+.set_attr<nnvm::FListInputNames>("FListInputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"queries_keys_values"};
+})
+.set_attr<nnvm::FListOutputNames>("FListOutputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"output"};
+})
+.set_attr<mxnet::FInferShape>("FInferShape", InterleavedMatMulSelfAttQKShape)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
+.set_attr<nnvm::FGradient>("FGradient",
+  ElemwiseGradUseIn{"_backward_interleaved_matmul_selfatt_qk"})
+.add_argument("queries_keys_values", "NDArray-or-Symbol", "Interleaved queries, keys and values")
+.add_arguments(InterleavedMatMulParam::__FIELDS__());
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_selfatt_qk)
+.set_num_inputs(2)
+.set_num_outputs(1)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>);
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_selfatt_valatt)
+.describe(R"code(Compute the matrix multiplication between the projections of
+values and the attention weights in multihead attention use as self attention.
+
+the inputs must be a tensor of interleaved projections
+of queries, keys and values following the layout:
+(seq_length, batch_size, num_heads * head_dim * 3)
+
+and the attention weights following the layout:
+(batch_size, seq_length, seq_length)
+
+the equivalent code would be:
+tmp = mx.nd.reshape(queries_keys_values, shape=(0, 0, num_heads, 3, -1))
+v_proj = mx.nd.transpose(tmp[:,:,:,2,:], axes=(1, 2, 0, 3))
+v_proj = mx.nd.reshape(v_proj, shape=(-1, 0, 0), reverse=True)
+output = mx.nd.batch_dot(attention, v_proj, transpose_b=True)
+output = mx.nd.reshape(output, shape=(-1, num_heads, 0, 0), reverse=True)
+output = mx.nd.transpose(output, axes=(0, 2, 1, 3))
+output = mx.nd.reshape(output, shape=(0, 0, -1))
+
+This Op is GPU only
+)code" ADD_FILELINE)
+.set_num_inputs(2)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>)
+.set_attr<nnvm::FListInputNames>("FListInputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"queries_keys_values", "attention"};
+})
+.set_attr<nnvm::FListOutputNames>("FListOutputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"output"};
+})
+.set_attr<mxnet::FInferShape>("FInferShape", InterleavedMatMulSelfAttValAttShape)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>)
+.set_attr<nnvm::FGradient>("FGradient",
+  ElemwiseGradUseIn{"_backward_interleaved_matmul_selfatt_valatt"})
+.add_argument("queries_keys_values", "NDArray-or-Symbol", "Queries, keys and values interleaved")
+.add_argument("attention", "NDArray-or-Symbol", "Attention maps")
+.add_arguments(InterleavedMatMulParam::__FIELDS__());
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_selfatt_valatt)
+.set_num_inputs(3)
+.set_num_outputs(2)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>);
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_encdec_qk)
+.describe(R"code(Compute the matrix multiplication between the projections of
+queries and keys in multihead attention use as encoder-decoder.
+
+the inputs must be a tensor of projections of queries following the layout:
+(seq_length, batch_size, num_heads * head_dim)
+
+and a tensor of interleaved projections of values and keys following the layout:
+(seq_length, batch_size, num_heads * head_dim * 2)
+
+the equivalent code would be:
+q_proj = mx.nd.transpose(queries, axes=(1, 2, 0, 3))
+q_proj = mx.nd.reshape(q_proj, shape=(-1, 0, 0), reverse=True)
+q_proj = mx.nd.contrib.div_sqrt_dim(q_proj)
+tmp = mx.nd.reshape(keys_values, shape=(0, 0, num_heads, 2, -1))
+k_proj = mx.nd.transpose(tmp[:,:,:,0,:], axes=(1, 2, 0, 3))
+k_proj = mx.nd.reshap(k_proj, shape=(-1, 0, 0), reverse=True)
+output = mx.nd.batch_dot(q_proj, k_proj, transpose_b=True)
+
+This Op is GPU only
+)code" ADD_FILELINE)
+.set_num_inputs(2)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>)
+.set_attr<nnvm::FListInputNames>("FListInputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"queries", "keys_values"};
+})
+.set_attr<nnvm::FListOutputNames>("FListOutputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"output"};
+})
+.set_attr<mxnet::FInferShape>("FInferShape", InterleavedMatMulEncDecQKShape)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>)
+.set_attr<nnvm::FGradient>("FGradient",
+    ElemwiseGradUseIn{"_backward_interleaved_matmul_encdec_qk"})
+.add_argument("queries", "NDArray-or-Symbol", "Queries")
+.add_argument("keys_values", "NDArray-or-Symbol", "Keys and values interleaved")
+.add_arguments(InterleavedMatMulParam::__FIELDS__());
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_encdec_qk)
+.set_num_inputs(3)
+.set_num_outputs(2)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>);
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_encdec_valatt)
+.describe(R"code(Compute the matrix multiplication between the projections of
+values and the attention weights in multihead attention use as encoder-decoder.
+
+the inputs must be a tensor of interleaved projections of
+keys and values following the layout:
+(seq_length, batch_size, num_heads * head_dim * 2)
+
+and the attention weights following the layout:
+(batch_size, seq_length, seq_length)
+
+the equivalent code would be:
+
+tmp = mx.nd.reshape(queries_keys_values, shape=(0, 0, num_heads, 3, -1))
+v_proj = mx.nd.transpose(tmp[:,:,:,1,:], axes=(1, 2, 0, 3))
+v_proj = mx.nd.reshape(v_proj, shape=(-1, 0, 0), reverse=True)
+output = mx.nd.batch_dot(attention, v_proj, transpose_b=True)
+output = mx.nd.reshape(output, shape=(-1, num_heads, 0, 0), reverse=True)
+output = mx.nd.transpose(output, axes=(0, 2, 1, 3))
+output = mx.nd.reshape(output, shape=(0, 0, -1))
+
+This Op is GPU only
+)code" ADD_FILELINE)
+.set_num_inputs(2)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>)
+.set_attr<nnvm::FListInputNames>("FListInputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"keys_values", "attention"};
+})
+.set_attr<nnvm::FListOutputNames>("FListOutputNames", [](const NodeAttrs& attrs) {
+  return std::vector<std::string>{"output"};
+})
+.set_attr<mxnet::FInferShape>("FInferShape", InterleavedMatMulEncDecValAttShape)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<2, 1>)
+.set_attr<nnvm::FGradient>("FGradient",
+    ElemwiseGradUseIn{"_backward_interleaved_matmul_encdec_valatt"})
+.add_argument("keys_values", "NDArray-or-Symbol", "Keys and values interleaved")
+.add_argument("attention", "NDArray-or-Symbol", "Attention maps")
+.add_arguments(InterleavedMatMulParam::__FIELDS__());
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_encdec_valatt)
+.set_num_inputs(3)
+.set_num_outputs(2)
+.set_attr<nnvm::TIsBackward>("TIsBackward", true)
+.set_attr_parser(ParamParser<InterleavedMatMulParam>);
+
+
 // relu
 MXNET_OPERATOR_REGISTER_UNARY(_contrib_div_sqrt_dim)
 .describe(R"code(Rescale the input by the square root of the channel dimension.
diff --git a/src/operator/contrib/transformer.cu b/src/operator/contrib/transformer.cu
index 6ed073db6011..e152669478dd 100644
--- a/src/operator/contrib/transformer.cu
+++ b/src/operator/contrib/transformer.cu
@@ -22,12 +22,572 @@
  * \file transformer.cu
  * \brief GPU implementation of the operators used in Transformer
  */
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
 #include <mxnet/base.h>
 #include "./transformer-inl.h"
+#include "../../common/cuda_utils.h"
 
 namespace mxnet {
 namespace op {
 
+// Approach in gemm_switch_fp32accum is coming from MLPerf v0.6 submission repository from NVIDIA
+// by https://github.com/kevinstephano
+template<typename DType>
+void CublasStridedBatchedGemm(mshadow::Stream<gpu>* s, bool transA, bool transB,
+                              int32_t m, int32_t n, int32_t k,
+                              float alpha, const DType* a, int32_t lda, int32_t strideA,
+                              const DType *b, int32_t ldb, int32_t strideB, float beta,
+                              DType *c, int32_t ldc, int32_t strideC, int32_t batchCount,
+                              cublasGemmAlgo_t algo = CUBLAS_GEMM_DEFAULT_TENSOR_OP) {
+#if CUDA_VERSION >= 9010
+  using namespace mxnet::common::cuda;
+  CHECK_EQ(s->blas_handle_ownership_, mshadow::Stream<gpu>::OwnHandle)
+      << "Must init CuBLAS handle in stream";
+
+  cublasHandle_t blas_handle = mshadow::Stream<gpu>::GetBlasHandle(s);
+  auto err = CUBLAS_STATUS_SUCCESS;
+  // TODO(cfujitsang): handle computation_precision
+  err = cublasGemmStridedBatchedEx(
+      blas_handle, CublasTransposeOp(transA), CublasTransposeOp(transB),
+      static_cast<int>(m), static_cast<int>(n), static_cast<int>(k),
+      reinterpret_cast<void*>(&alpha),
+      a, CublasType<DType>::kCudaFlag, static_cast<int>(lda), strideA,
+      b, CublasType<DType>::kCudaFlag, static_cast<int>(ldb), strideB,
+      reinterpret_cast<void*>(&beta),
+      c, CublasType<DType>::kCudaFlag, static_cast<int>(ldc), strideC,
+      static_cast<int>(batchCount), CUDA_R_32F, algo);
+  CHECK_EQ(err, CUBLAS_STATUS_SUCCESS) << "Cublas gemmEx fail.";
+#else
+  LOG(FATAL) << "Not implemented with CUDA < 9.1";
+#endif
+}
+
+template<typename DType>
+void gemm_switch_fp32accum(mshadow::Stream<gpu>* s, bool transA, bool transB,
+                           int32_t m, int32_t n, int32_t k,
+                           float alpha, const DType *a, int32_t lda,
+                           int32_t strideA, const DType *b, int32_t ldb,
+                           int32_t strideB, float beta, DType *c, int32_t ldc,
+                           int32_t strideC, int32_t batchCount) {
+  cudaStream_t stream = mshadow::Stream<gpu>::GetStream(s);
+  if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) {
+    CublasStridedBatchedGemm(s, transA, transB, m, n, k, alpha, a, lda, strideA, b, ldb,
+      strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);
+  } else {
+    CublasStridedBatchedGemm(s, transA, transB, m, n, k, alpha, a, lda, strideA, b, ldb,
+      strideB, beta, c, ldc, strideC, batchCount);
+  }
+  CHECK_CUDA_ERROR("Error at InterleavedMatMul");
+}
+
+// TODO(cfujitsang): use scale as optional ?
+void InterleavedMatMulSelfAttQKGPU(const nnvm::NodeAttrs& attrs,
+                                   const OpContext &ctx,
+                                   const std::vector<TBlob> &inputs,
+                                   const std::vector<OpReqType> &req,
+                                   const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* queries_keys_values = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* output = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t qkv_seq_len    = inputs[0].shape_[0];
+    const int32_t sequences      = inputs[0].shape_[1];
+    const int32_t output_lin_dim = inputs[0].shape_[2];
+    const int32_t embed_dim      = output_lin_dim / 3;
+    const int32_t head_dim       = embed_dim / params.heads;
+    const int32_t attn_batches   = params.heads * sequences;
+    const int32_t lead_dim       = attn_batches * 3 * head_dim;
+    const int32_t batch_stride   = 3 * head_dim;
+    const float beta             = req[0] == kAddTo ? 1.f : 0.f;
+    const float scale            = 1.0 / sqrt(static_cast<float>(head_dim));
+
+    if (req[0] == kNullOp)
+      return;
+
+    gemm_switch_fp32accum(s,
+                          true,
+                          false,
+                          qkv_seq_len,
+                          qkv_seq_len,
+                          head_dim,
+                          scale,
+                          queries_keys_values + head_dim,
+                          lead_dim,
+                          batch_stride,
+                          queries_keys_values,
+                          lead_dim,
+                          batch_stride,
+                          beta,
+                          output,
+                          qkv_seq_len,
+                          qkv_seq_len * qkv_seq_len,
+                          attn_batches);
+  })
+}
+
+void BackwardInterleavedMatMulSelfAttQKGPU(const nnvm::NodeAttrs& attrs,
+                                           const OpContext &ctx,
+                                           const std::vector<TBlob> &inputs,
+                                           const std::vector<OpReqType> &req,
+                                           const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* output_grads        = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* queries_keys_values = inputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* queries_keys_values_grads = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t qkv_seq_len    = inputs[1].shape_[0];
+    const int32_t sequences      = inputs[1].shape_[1];
+    const int32_t output_lin_dim = inputs[1].shape_[2];
+    const int32_t embed_dim      = output_lin_dim / 3;
+    const int32_t head_dim       = embed_dim / params.heads;
+    const int32_t attn_batches   = params.heads * sequences;
+    const int32_t lead_dim       = attn_batches * 3 * head_dim;
+    const int32_t batch_stride   = 3 * head_dim;
+    const float scale            = 1.0 / sqrt(static_cast<float>(head_dim));
+    const float beta = req[0] == kAddTo ? 1.f : 0.f;
+
+    if (req[0] == kNullOp)
+      return;
+
+    if (req[0] == kWriteTo) {
+      cudaMemsetAsync(queries_keys_values_grads, 0, outputs[0].shape_.Size() * sizeof(DType),
+                      mshadow::Stream<gpu>::GetStream(s));
+    }
+
+    gemm_switch_fp32accum(s,
+                          false,
+                          false,
+                          head_dim,
+                          qkv_seq_len,
+                          qkv_seq_len,
+                          scale,
+                          queries_keys_values + head_dim,
+                          lead_dim,
+                          batch_stride,
+                          output_grads,
+                          qkv_seq_len,
+                          qkv_seq_len * qkv_seq_len,
+                          beta,
+                          queries_keys_values_grads,
+                          lead_dim,
+                          batch_stride,
+                          attn_batches);
+    gemm_switch_fp32accum(s,
+                          false,
+                          true,
+                          head_dim,
+                          qkv_seq_len,
+                          qkv_seq_len,
+                          scale,
+                          queries_keys_values,
+                          lead_dim,
+                          batch_stride,
+                          output_grads,
+                          qkv_seq_len,
+                          qkv_seq_len * qkv_seq_len,
+                          beta,
+                          queries_keys_values_grads + head_dim,
+                          lead_dim,
+                          batch_stride,
+                          attn_batches);
+  })
+}
+
+void InterleavedMatMulSelfAttValAttGPU(const nnvm::NodeAttrs& attrs,
+                                       const OpContext &ctx,
+                                       const std::vector<TBlob> &inputs,
+                                       const std::vector<OpReqType> &req,
+                                       const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* queries_keys_values = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* attention_maps      = inputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* output                    = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t qkv_seq_len    = inputs[0].shape_[0];
+    const int32_t sequences      = inputs[0].shape_[1];
+    const int32_t output_lin_dim = inputs[0].shape_[2];
+    const int32_t embed_dim      = output_lin_dim / 3;
+    const int32_t head_dim       = embed_dim / params.heads;
+    const int32_t attn_batches   = params.heads * sequences;
+    const int32_t lead_dim       = attn_batches * 3 * head_dim;
+    const int32_t batch_stride   = 3 * head_dim;
+    const float alpha            = 1.f;
+    const float beta             = req[0] == kAddTo ? 1.f : 0.f;
+
+    if (req[0] == kNullOp)
+      return;
+
+    gemm_switch_fp32accum(s,
+                          false,
+                          false,
+                          head_dim,
+                          qkv_seq_len,
+                          qkv_seq_len,
+                          alpha,
+                          queries_keys_values + 2 * head_dim,
+                          lead_dim,
+                          batch_stride,
+                          attention_maps,
+                          qkv_seq_len,
+                          qkv_seq_len * qkv_seq_len,
+                          beta,
+                          output,
+                          head_dim * attn_batches,
+                          head_dim,
+                          attn_batches);
+  })
+}
+
+void BackwardInterleavedMatMulSelfAttValAttGPU(const nnvm::NodeAttrs& attrs,
+                                               const OpContext &ctx,
+                                               const std::vector<TBlob> &inputs,
+                                               const std::vector<OpReqType> &req,
+                                               const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* output_grads              = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* queries_keys_values       = inputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* attention_maps            = inputs[2].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* queries_keys_values_grads       = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* attention_maps_grads            = outputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t qkv_seq_len    = inputs[1].shape_[0];
+    const int32_t sequences      = inputs[1].shape_[1];
+    const int32_t output_lin_dim = inputs[1].shape_[2];
+    const int32_t embed_dim      = output_lin_dim / 3;
+    const int32_t head_dim       = embed_dim / params.heads;
+    const int32_t attn_batches   = params.heads * sequences;
+    const int32_t lead_dim       = attn_batches * 3 * head_dim;
+    const int32_t batch_stride   = 3 * head_dim;
+    const float alpha            = 1.f;
+    if (req[0] != kNullOp) {
+      if (req[0] == kWriteTo) {
+        cudaMemsetAsync(queries_keys_values_grads, 0, outputs[0].shape_.Size() * sizeof(DType),
+                        mshadow::Stream<gpu>::GetStream(s));
+      }
+      const float beta = req[0] == kAddTo ? 1.f : 0.f;
+      gemm_switch_fp32accum(s,
+                            false,
+                            true,
+                            head_dim,
+                            qkv_seq_len,
+                            qkv_seq_len,
+                            alpha,
+                            output_grads,
+                            head_dim * attn_batches,
+                            head_dim,
+                            attention_maps,
+                            qkv_seq_len,
+                            qkv_seq_len * qkv_seq_len,
+                            beta,
+                            queries_keys_values_grads + 2 * head_dim,
+                            lead_dim,
+                            batch_stride,
+                            attn_batches);
+    }
+    if (req[1] != kNullOp) {
+      const float beta = req[1] == kAddTo ? 1.f : 0.f;
+      gemm_switch_fp32accum(s,
+                            true,
+                            false,
+                            qkv_seq_len,
+                            qkv_seq_len,
+                            head_dim,
+                            alpha,
+                            queries_keys_values + 2 * head_dim,
+                            lead_dim,
+                            batch_stride,
+                            output_grads,
+                            head_dim * attn_batches,
+                            head_dim,
+                            beta,
+                            attention_maps_grads,
+                            qkv_seq_len,
+                            qkv_seq_len * qkv_seq_len,
+                            attn_batches);
+    }
+  })
+}
+
+
+void InterleavedMatMulEncDecQKGPU(const nnvm::NodeAttrs& attrs,
+                                  const OpContext &ctx,
+                                  const std::vector<TBlob> &inputs,
+                                  const std::vector<OpReqType> &req,
+                                  const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* queries     = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* keys_values = inputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* output            = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t q_seq_len         = inputs[0].shape_[0];
+    const int32_t sequences         = inputs[0].shape_[1];
+    const int32_t output_lin_q_dim  = inputs[0].shape_[2];
+    const int32_t kv_seq_len        = inputs[1].shape_[0];
+    const int32_t output_lin_kv_dim = inputs[1].shape_[2];
+    const int32_t embed_dim         = output_lin_q_dim;
+    const int32_t head_dim          = embed_dim / params.heads;
+    const int32_t attn_batches      = params.heads * sequences;
+    const int32_t lead_dim_q        = attn_batches * head_dim;
+    const int32_t lead_dim_kv       = attn_batches * 2 * head_dim;
+    const int32_t batch_stride_q    = head_dim;
+    const int32_t batch_stride_kv   = head_dim * 2;
+    const float beta                = req[0] == kAddTo ? 1.f : 0.f;
+    const float scale               = 1.f / sqrt(static_cast<float>(head_dim));
+
+    if (req[0] == kNullOp)
+      return;
+
+    gemm_switch_fp32accum(s,
+                          true,
+                          false,
+                          kv_seq_len,
+                          q_seq_len,
+                          head_dim,
+                          scale,
+                          keys_values,
+                          lead_dim_kv,
+                          batch_stride_kv,
+                          queries,
+                          lead_dim_q,
+                          batch_stride_q,
+                          beta,
+                          output,
+                          kv_seq_len,
+                          kv_seq_len * q_seq_len,
+                          attn_batches);
+  })
+}
+
+void BackwardInterleavedMatMulEncDecQKGPU(const nnvm::NodeAttrs& attrs,
+                                          const OpContext &ctx,
+                                          const std::vector<TBlob> &inputs,
+                                          const std::vector<OpReqType> &req,
+                                          const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* output_grads = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* queries       = inputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* keys_values   = inputs[2].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* queries_grads       = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* keys_values_grads   = outputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t q_seq_len         = inputs[1].shape_[0];
+    const int32_t sequences         = inputs[1].shape_[1];
+    const int32_t output_lin_q_dim  = inputs[1].shape_[2];
+    const int32_t kv_seq_len        = inputs[2].shape_[0];
+    const int32_t output_lin_kv_dim = inputs[2].shape_[2];
+    const int32_t embed_dim         = output_lin_q_dim;
+    const int32_t head_dim          = embed_dim / params.heads;
+    const int32_t attn_batches      = params.heads * sequences;
+    const int32_t lead_dim_q        = attn_batches * head_dim;
+    const int32_t lead_dim_kv       = attn_batches * 2 * head_dim;
+    const int32_t batch_stride_q    = head_dim;
+    const int32_t batch_stride_kv   = head_dim * 2;
+    const float scale               = 1.f / sqrt(static_cast<float>(head_dim));
+
+    if (req[0] != kNullOp) {
+      const float beta = req[0] == kAddTo ? 1.f : 0.f;
+      gemm_switch_fp32accum(s,
+                            false,
+                            false,
+                            head_dim,
+                            q_seq_len,
+                            kv_seq_len,
+                            scale,
+                            keys_values,
+                            lead_dim_kv,
+                            batch_stride_kv,
+                            output_grads,
+                            kv_seq_len,
+                            kv_seq_len * q_seq_len,
+                            beta,
+                            queries_grads,
+                            lead_dim_q,
+                            batch_stride_q,
+                            attn_batches);
+    }
+    if (req[1] != kNullOp) {
+      if (req[1] == kWriteTo) {
+        cudaMemsetAsync(keys_values_grads, 0, outputs[1].shape_.Size() * sizeof(DType),
+                        mshadow::Stream<gpu>::GetStream(s));
+      }
+      const float beta = req[1] == kAddTo ? 1.f : 0.f;
+      gemm_switch_fp32accum(s,
+                            false,
+                            true,
+                            head_dim,
+                            kv_seq_len,
+                            q_seq_len,
+                            scale,
+                            queries,
+                            lead_dim_q,
+                            batch_stride_q,
+                            output_grads,
+                            kv_seq_len,
+                            kv_seq_len * q_seq_len,
+                            beta,
+                            keys_values_grads,
+                            lead_dim_kv,
+                            batch_stride_kv,
+                            attn_batches);
+    }
+  })
+}
+
+void InterleavedMatMulEncDecValAttGPU(const nnvm::NodeAttrs& attrs,
+                                      const OpContext &ctx,
+                                      const std::vector<TBlob> &inputs,
+                                      const std::vector<OpReqType> &req,
+                                      const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* keys_values    = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* attention_maps = inputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* output               = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t kv_seq_len        = inputs[0].shape_[0];
+    const int32_t sequences         = inputs[0].shape_[1];
+    const int32_t output_lin_kv_dim = inputs[0].shape_[2];
+    const int32_t attn_batches      = inputs[1].shape_[0];
+    const int32_t q_seq_len         = inputs[1].shape_[1];
+    const int32_t embed_dim         = output_lin_kv_dim / 2;
+    int32_t head_dim                = embed_dim / params.heads;
+    const int32_t lead_dim_kv       = attn_batches * head_dim * 2;
+    const int32_t batch_stride_kv   = 2 * head_dim;
+    const float alpha               = 1.f;
+    const float beta                = req[0] == kAddTo ? 1.f : 0.f;
+
+    if (req[0] == kNullOp)
+      return;
+
+    gemm_switch_fp32accum(s,
+                          false,
+                          false,
+                          head_dim,
+                          q_seq_len,
+                          kv_seq_len,
+                          alpha,
+                          keys_values + head_dim,
+                          lead_dim_kv,
+                          batch_stride_kv,
+                          attention_maps,
+                          kv_seq_len,
+                          kv_seq_len * q_seq_len,
+                          beta,
+                          output,
+                          head_dim * attn_batches,
+                          head_dim,
+                          attn_batches);
+  })
+}
+
+void BackwardInterleavedMatMulEncDecValAttGPU(const nnvm::NodeAttrs& attrs,
+                                              const OpContext &ctx,
+                                              const std::vector<TBlob> &inputs,
+                                              const std::vector<OpReqType> &req,
+                                              const std::vector<TBlob> &outputs) {
+  const auto& params = nnvm::get<InterleavedMatMulParam>(attrs.parsed);
+  mshadow::Stream<gpu>* s = ctx.get_stream<gpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    const DType* output_grads   = inputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* keys_values    = inputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    const DType* attention_maps = inputs[2].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* keys_values_grads    = outputs[0].FlatTo2D<gpu, DType>(s).dptr_;
+    DType* attention_maps_grads = outputs[1].FlatTo2D<gpu, DType>(s).dptr_;
+    const int32_t kv_seq_len        = inputs[1].shape_[0];
+    const int32_t sequences         = inputs[1].shape_[1];
+    const int32_t output_lin_kv_dim = inputs[1].shape_[2];
+    const int32_t attn_batches      = inputs[2].shape_[0];
+    const int32_t q_seq_len         = inputs[2].shape_[1];
+    const int32_t embed_dim         = output_lin_kv_dim / 2;
+    int32_t head_dim                = embed_dim / params.heads;
+    const int32_t lead_dim_kv       = attn_batches * head_dim * 2;
+    const int32_t batch_stride_kv   = 2 * head_dim;
+    const float alpha               = 1.f;
+
+    if (req[0] != kNullOp) {
+      if (req[0] == kWriteTo) {
+        cudaMemsetAsync(keys_values_grads, 0, outputs[0].shape_.Size() * sizeof(DType),
+                        mshadow::Stream<gpu>::GetStream(s));
+      }
+      const float beta = req[0] == kAddTo ? 1.f : 0.f;
+      gemm_switch_fp32accum(s,
+                            false,
+                            true,
+                            head_dim,
+                            kv_seq_len,
+                            q_seq_len,
+                            alpha,
+                            output_grads,
+                            head_dim * attn_batches,
+                            head_dim,
+                            attention_maps,
+                            kv_seq_len,
+                            kv_seq_len * q_seq_len,
+                            beta,
+                            keys_values_grads + head_dim,
+                            lead_dim_kv,
+                            batch_stride_kv,
+                            attn_batches);
+    }
+    if (req[1] != kNullOp) {
+      const float beta = req[1] == kAddTo ? 1.f : 0.f;
+      gemm_switch_fp32accum(s,
+                            true,
+                            false,
+                            kv_seq_len,
+                            q_seq_len,
+                            head_dim,
+                            alpha,
+                            keys_values + head_dim,
+                            lead_dim_kv,
+                            batch_stride_kv,
+                            output_grads,
+                            head_dim * attn_batches,
+                            head_dim,
+                            beta,
+                            attention_maps_grads,
+                            kv_seq_len,
+                            kv_seq_len * q_seq_len,
+                            attn_batches);
+    }
+  })
+}
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_selfatt_qk)
+.set_attr<FCompute>("FCompute<gpu>", InterleavedMatMulSelfAttQKGPU);
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_selfatt_valatt)
+.set_attr<FCompute>("FCompute<gpu>", InterleavedMatMulSelfAttValAttGPU);
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_encdec_qk)
+.set_attr<FCompute>("FCompute<gpu>", InterleavedMatMulEncDecQKGPU);
+
+NNVM_REGISTER_OP(_contrib_interleaved_matmul_encdec_valatt)
+.set_attr<FCompute>("FCompute<gpu>", InterleavedMatMulEncDecValAttGPU);
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_selfatt_qk)
+.set_attr<FCompute>("FCompute<gpu>", BackwardInterleavedMatMulSelfAttQKGPU);
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_selfatt_valatt)
+.set_attr<FCompute>("FCompute<gpu>", BackwardInterleavedMatMulSelfAttValAttGPU);
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_encdec_qk)
+.set_attr<FCompute>("FCompute<gpu>", BackwardInterleavedMatMulEncDecQKGPU);
+
+NNVM_REGISTER_OP(_backward_interleaved_matmul_encdec_valatt)
+.set_attr<FCompute>("FCompute<gpu>", BackwardInterleavedMatMulEncDecValAttGPU);
+
 // relu
 NNVM_REGISTER_OP(_contrib_div_sqrt_dim)
 .set_attr<FCompute>("FCompute<gpu>", DivSqrtDimForward_<gpu>);
diff --git a/tests/python/gpu/test_operator_gpu.py b/tests/python/gpu/test_operator_gpu.py
index 0f1cd93755c3..a8e70de18d1b 100644
--- a/tests/python/gpu/test_operator_gpu.py
+++ b/tests/python/gpu/test_operator_gpu.py
@@ -2493,13 +2493,327 @@ def test_arange_like_dtype():
         x = mx.sym.Variable('x', dtype=t)
         y = mx.sym.reshape(x, shape=(0, 0, -1))
         z = mx.sym.contrib.arange_like(y, axis=-1)
-    
+
         mod = z.simple_bind(ctx=mx.gpu(0), x=(3, 4, 5, 6), grad_req='null')
         mod.arg_arrays[0][:] = np.random.normal(size=mod.arg_arrays[0].shape).astype(t)
         out = mod.forward(is_train=False)
         for v in out:
             assert v.dtype == t
 
+@with_seed()
+def check_multihead_attention_selfatt(dtype):
+    def convert_weight(F, q_weight, k_weight, v_weight, num_heads):
+        q_weight = F.reshape(q_weight, shape=(num_heads, -1, 0), reverse=True)
+        k_weight = F.reshape(k_weight, shape=(num_heads, -1, 0), reverse=True)
+        v_weight = F.reshape(v_weight, shape=(num_heads, -1, 0), reverse=True)
+        all_weights = F.concat(q_weight, k_weight, v_weight, dim=-2)
+        all_weights = F.reshape(all_weights, shape=(-1, 0), reverse=True)
+        return all_weights
+
+    def convert_bias(F, q_bias, k_bias, v_bias, num_heads):
+        q_bias = F.reshape(q_bias, shape=(num_heads, -1))
+        k_bias = F.reshape(k_bias, shape=(num_heads, -1))
+        v_bias = F.reshape(v_bias, shape=(num_heads, -1))
+        all_bias = F.stack(q_bias, k_bias, v_bias, axis=1)
+        all_bias = F.reshape(all_bias, shape=(-1,))
+        return all_bias
+
+    batch_size = 2
+    qkv_length = 7  # length of a sequence
+    qkv_dim = 9     # dimension of encoding
+    num_heads = 3   # number of attention head
+    head_dim = 5    # head size
+    out_dim = 13 * num_heads
+    qkv_units = num_heads * head_dim
+
+    arg_params = {
+        'qkv': mx.nd.array(np.random.rand(*(batch_size, qkv_length, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'q_weight': mx.nd.array(np.random.rand(*(qkv_units, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'k_weight': mx.nd.array(np.random.rand(*(qkv_units, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'v_weight': mx.nd.array(np.random.rand(*(qkv_units, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'q_bias': mx.nd.array(np.random.rand(*(qkv_units,)).astype(dtype) * 0.1, dtype=dtype),
+        'k_bias': mx.nd.array(np.random.rand(*(qkv_units,)).astype(dtype) * 0.1, dtype=dtype),
+        'v_bias': mx.nd.array(np.random.rand(*(qkv_units,)).astype(dtype) * 0.1, dtype=dtype),
+        'out_weight': mx.nd.array(np.random.rand(*(out_dim, qkv_units)).astype(dtype) * 0.1, dtype=dtype),
+        'out_bias': mx.nd.array(np.random.rand(*(out_dim,)).astype(dtype) * 0.1, dtype=dtype),
+        }
+
+    qkv = mx.sym.Variable('qkv')
+    sonde = mx.sym.Variable('sonde')
+    q_weight = mx.sym.Variable('q_weight')
+    k_weight = mx.sym.Variable('k_weight')
+    v_weight = mx.sym.Variable('v_weight')
+    q_bias = mx.sym.Variable('q_bias')
+    k_bias = mx.sym.Variable('k_bias')
+    v_bias = mx.sym.Variable('v_bias')
+    out_weight = mx.sym.Variable('out_weight')
+    out_bias = mx.sym.Variable('out_bias')
+    qkv_weight = convert_weight(mx.sym, q_weight, k_weight, v_weight, num_heads)
+    qkv_bias = convert_bias(mx.sym, q_bias, k_bias, v_bias, num_heads)
+    qkv = mx.sym.transpose(qkv, axes=(1, 0, 2))
+    qkv_proj = mx.sym.FullyConnected(qkv, weight=qkv_weight, bias=qkv_bias, flatten=False,
+                                     num_hidden=qkv_units * 3, no_bias=False)
+    att_score = mx.sym.contrib.interleaved_matmul_selfatt_qk(
+            qkv_proj, heads=num_heads)
+    att_score = att_score + sonde
+    weighted_value = mx.sym.contrib.interleaved_matmul_selfatt_valatt(
+            qkv_proj, att_score, heads=num_heads)
+    output = mx.sym.FullyConnected(weighted_value, weight=out_weight, bias=out_bias, flatten=False,
+                                   num_hidden=out_dim, no_bias=False)
+    output = mx.sym.transpose(output, axes=(1, 0, 2))
+    output = mx.sym.Group([output, att_score])
+    executor = output.simple_bind(ctx=mx.gpu(0),
+                                  qkv=(batch_size, qkv_length, qkv_dim),
+                                  q_weight=(qkv_units, qkv_dim),
+                                  q_bias=(qkv_units,),
+                                  k_weight=(qkv_units, qkv_dim),
+                                  k_bias=(qkv_units,),
+                                  v_weight=(qkv_units, qkv_dim),
+                                  v_bias=(qkv_units,),
+                                  type_dict={'qkv': dtype,
+                                             'q_weight': dtype,
+                                             'k_weight': dtype,
+                                             'v_weight': dtype,
+                                             'q_bias': dtype,
+                                             'k_bias': dtype,
+                                             'v_bias': dtype,
+                                             'sonde': dtype},
+                                  grad_req='write', force_rebind=True)
+    output_shape = executor.outputs[0].shape
+    output_grads = np.random.rand(*output_shape).astype(dtype) * 0.1
+    executor.copy_params_from(arg_params, {})
+    executor.arg_dict['sonde'][:] = 0.
+    executor.arg_dict['sonde'].wait_to_read()
+    executor.forward(is_train=True)
+    output_opti = executor.outputs[0].asnumpy()
+    att_score_opti = executor.outputs[1].asnumpy()
+    executor.backward([mx.nd.array(output_grads, dtype=dtype),
+                       mx.nd.zeros(att_score_opti.shape, dtype=dtype)])
+    grads_opti = {k: v.asnumpy() for k, v in executor.grad_dict.items()}
+    qkv = mx.sym.Variable('qkv')
+    sonde = mx.sym.Variable('sonde')
+    q_weight = mx.sym.Variable('q_weight')
+    k_weight = mx.sym.Variable('k_weight')
+    v_weight = mx.sym.Variable('v_weight')
+    q_bias = mx.sym.Variable('q_bias')
+    k_bias = mx.sym.Variable('k_bias')
+    v_bias = mx.sym.Variable('v_bias')
+    out_weight = mx.sym.Variable('out_weight')
+    out_bias = mx.sym.Variable('out_bias')
+
+    q = mx.sym.FullyConnected(qkv, weight=q_weight, bias=q_bias, flatten=False,
+                              num_hidden=qkv_units, no_bias=False)
+    k = mx.sym.FullyConnected(qkv, weight=k_weight, bias=k_bias, flatten=False,
+                              num_hidden=qkv_units, no_bias=False)
+    v = mx.sym.FullyConnected(qkv, weight=v_weight, bias=v_bias, flatten=False,
+                              num_hidden=qkv_units, no_bias=False)
+    q = mx.sym.reshape(q, shape=(0, 0, num_heads, -1))
+    q = mx.sym.transpose(q, axes=(0, 2, 1, 3))
+    q = mx.sym.reshape(q, shape=(-1, 0, 0), reverse=True)
+    k = mx.sym.reshape(k, shape=(0, 0, num_heads, -1))
+    k = mx.sym.transpose(k, axes=(0, 2, 1, 3))
+    k = mx.sym.reshape(k, shape=(-1, 0, 0), reverse=True)
+    q = mx.sym.contrib.div_sqrt_dim(q)
+    att_score = mx.sym.batch_dot(q, k, transpose_b=True)
+    att_score = att_score + sonde
+    v = mx.sym.reshape(v, shape=(0, 0, num_heads, -1))
+    v = mx.sym.transpose(v, axes=(0, 2, 1, 3))
+    v = mx.sym.reshape(v, shape=(-1, 0, 0), reverse=True)
+    weighted_value = mx.sym.batch_dot(att_score, v)
+    weighted_value = mx.sym.reshape(weighted_value, shape=(-1, num_heads, 0, 0),
+                                    reverse=True)
+    weighted_value = mx.sym.transpose(weighted_value, axes=(0, 2, 1, 3))
+    weighted_value = mx.sym.reshape(weighted_value, shape=(0, 0, -1))
+    output = mx.sym.FullyConnected(weighted_value, weight=out_weight, bias=out_bias, flatten=False,
+                                   num_hidden=out_dim, no_bias=False)
+    output = mx.sym.Group([output, att_score])
+    executor = output.simple_bind(ctx=mx.gpu(0),
+                                  qkv=(batch_size, qkv_length, qkv_dim),
+                                  type_dict={'qkv': dtype},
+                                  grad_req='write', force_rebind=True)
+    executor.copy_params_from(arg_params, {})
+    executor.arg_dict['sonde'][:] = 0.
+    executor.arg_dict['sonde'].wait_to_read()
+    executor.forward(is_train=True)
+    output_orig = executor.outputs[0].asnumpy()
+    att_score_orig = executor.outputs[1].asnumpy()
+    executor.backward([mx.nd.array(output_grads, dtype=dtype),
+                       mx.nd.zeros(att_score_orig.shape, dtype=dtype)])
+    grads_orig = {k : v.asnumpy() for k, v in executor.grad_dict.items()}
+    assert_allclose(att_score_orig, att_score_opti, rtol=1e-2, atol=1e-3)
+    assert_allclose(output_orig, output_opti, rtol=1e-2, atol=1e-3)
+
+    for k in grads_opti.keys():
+        assert(grads_orig[k].dtype == grads_opti[k].dtype)
+        assert(grads_orig[k].shape == grads_opti[k].shape)
+        assert_allclose(grads_orig[k], grads_opti[k], rtol=1e-2, atol=1e-3)
+
+def test_multihead_attention_selfatt():
+    for dtype in ['float16', 'float32']:
+        check_multihead_attention_selfatt(dtype=dtype)
+
+def check_multihead_attention_encdec(dtype):
+    def convert_weight(F, k_weight, v_weight, num_heads):
+        k_weight = F.reshape(k_weight, shape=(num_heads, -1, 0), reverse=True)
+        v_weight = F.reshape(v_weight, shape=(num_heads, -1, 0), reverse=True)
+        all_weights = F.concat(k_weight, v_weight, dim=-2)
+        all_weights = F.reshape(all_weights, shape=(-1, 0), reverse=True)
+        return all_weights
+
+    def convert_bias(F, k_bias, v_bias, num_heads):
+        k_bias = F.reshape(k_bias, shape=(num_heads, -1))
+        v_bias = F.reshape(v_bias, shape=(num_heads, -1))
+        all_bias = F.stack(k_bias, v_bias, axis=1)
+        all_bias = F.reshape(all_bias, shape=(-1,))
+        return all_bias
+
+    batch_size = 2
+    qkv_length = 7  # length of a sequence
+    qkv_dim = 9     # dimension of encoding
+    num_heads = 3   # number of attention head
+    head_dim = 5    # head size
+    out_dim = 13 * num_heads
+    qkv_units = num_heads * head_dim
+
+    arg_params = {
+        'q': mx.nd.array(np.random.rand(*(batch_size, qkv_length, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'kv': mx.nd.array(np.random.rand(*(batch_size, qkv_length, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'q_weight': mx.nd.array(np.random.rand(*(qkv_units, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'k_weight': mx.nd.array(np.random.rand(*(qkv_units, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'v_weight': mx.nd.array(np.random.rand(*(qkv_units, qkv_dim)).astype(dtype) * 0.1, dtype=dtype),
+        'q_bias': mx.nd.array(np.random.rand(*(qkv_units,)).astype(dtype) * 0.1, dtype=dtype),
+        'k_bias': mx.nd.array(np.random.rand(*(qkv_units,)).astype(dtype) * 0.1, dtype=dtype),
+        'v_bias': mx.nd.array(np.random.rand(*(qkv_units,)).astype(dtype) * 0.1, dtype=dtype),
+        'out_weight': mx.nd.array(np.random.rand(*(out_dim, qkv_units)).astype(dtype) * 0.1, dtype=dtype),
+        'out_bias': mx.nd.array(np.random.rand(*(out_dim,)).astype(dtype) * 0.1, dtype=dtype),
+        }
+
+    q = mx.sym.Variable('q')
+    kv = mx.sym.Variable('kv')
+    sonde = mx.sym.Variable('sonde')
+    q_weight = mx.sym.Variable('q_weight')
+    k_weight = mx.sym.Variable('k_weight')
+    v_weight = mx.sym.Variable('v_weight')
+    q_bias = mx.sym.Variable('q_bias')
+    k_bias = mx.sym.Variable('k_bias')
+    v_bias = mx.sym.Variable('v_bias')
+    out_weight = mx.sym.Variable('out_weight')
+    out_bias = mx.sym.Variable('out_bias')
+    kv_weight = convert_weight(mx.sym, k_weight, v_weight, num_heads)
+    kv_bias = convert_bias(mx.sym, k_bias, v_bias, num_heads)
+    kv = mx.sym.transpose(kv, axes=(1, 0, 2))
+    kv_proj = mx.sym.FullyConnected(kv, weight=kv_weight, bias=kv_bias, flatten=False,
+                                    num_hidden=qkv_units * 2, no_bias=False)
+    q = mx.sym.transpose(q, axes=(1, 0, 2))
+    q_proj = mx.sym.FullyConnected(q, weight=q_weight, bias=q_bias, flatten=False,
+                                   num_hidden=qkv_units, no_bias=False)
+    att_score = mx.sym.contrib.interleaved_matmul_encdec_qk(
+            q_proj, kv_proj, heads=num_heads) 
+    att_score = att_score + sonde
+    weighted_value = mx.sym.contrib.interleaved_matmul_encdec_valatt(
+            kv_proj, att_score, heads=num_heads)
+    output = mx.sym.FullyConnected(weighted_value, weight=out_weight, bias=out_bias, flatten=False,
+                                   num_hidden=out_dim, no_bias=False)
+    output = mx.sym.transpose(output, axes=(1, 0, 2))
+    output = mx.sym.Group([output, att_score])
+    executor = output.simple_bind(ctx=mx.gpu(0),
+                                  q=(batch_size, qkv_length, qkv_dim),
+                                  kv=(batch_size, qkv_length, qkv_dim),
+                                  q_weight=(qkv_units, qkv_dim),
+                                  q_bias=(qkv_units,),
+                                  k_weight=(qkv_units, qkv_dim),
+                                  k_bias=(qkv_units,),
+                                  v_weight=(qkv_units, qkv_dim),
+                                  v_bias=(qkv_units,),
+                                  out_weight=(out_dim, qkv_units),
+                                  out_bias=(out_dim,),
+                                  type_dict={'q': dtype,
+                                             'kv': dtype,
+                                             'q_weight': dtype,
+                                             'q_bias': dtype,
+                                             'k_weight': dtype,
+                                             'k_bias': dtype,
+                                             'v_weight': dtype,
+                                             'v_bias': dtype,
+                                             'out_weight': dtype,
+                                             'out_bias': dtype,
+                                              },
+                                  grad_req='write', force_rebind=True)
+    output_shape = executor.outputs[0].shape
+    output_grads = np.random.rand(*output_shape).astype(dtype) * 0.1
+    executor.copy_params_from(arg_params, {})
+    executor.arg_dict['sonde'][:] = 0.
+    executor.arg_dict['sonde'].wait_to_read()
+    executor.forward(is_train=True)
+    output_opti = executor.outputs[0].asnumpy()
+    att_score_opti = executor.outputs[1].asnumpy()
+    executor.backward([mx.nd.array(output_grads, dtype=dtype), mx.nd.zeros(att_score_opti.shape, dtype=dtype)])
+
+    grads_opti = {k: v.asnumpy() for k, v in executor.grad_dict.items()}
+
+    q = mx.sym.Variable('q')
+    kv = mx.sym.Variable('kv')
+    sonde = mx.sym.Variable('sonde')
+    q_weight = mx.sym.Variable('q_weight')
+    k_weight = mx.sym.Variable('k_weight')
+    v_weight = mx.sym.Variable('v_weight')
+    q_bias = mx.sym.Variable('q_bias')
+    k_bias = mx.sym.Variable('k_bias')
+    v_bias = mx.sym.Variable('v_bias')
+    out_weight = mx.sym.Variable('out_weight')
+    out_bias = mx.sym.Variable('out_bias')
+
+    q = mx.sym.FullyConnected(q, weight=q_weight, bias=q_bias, flatten=False,
+                              num_hidden=qkv_units, no_bias=False)
+    k = mx.sym.FullyConnected(kv, weight=k_weight, bias=k_bias, flatten=False,
+                              num_hidden=qkv_units, no_bias=False)
+    v = mx.sym.FullyConnected(kv, weight=v_weight, bias=v_bias, flatten=False,
+                              num_hidden=qkv_units, no_bias=False)
+    q = mx.sym.reshape(q, shape=(0, 0, num_heads, -1))
+    q = mx.sym.transpose(q, axes=(0, 2, 1, 3))
+    q = mx.sym.reshape(q, shape=(-1, 0, 0), reverse=True)
+    k = mx.sym.reshape(k, shape=(0, 0, num_heads, -1))
+    k = mx.sym.transpose(k, axes=(0, 2, 1, 3))
+    k = mx.sym.reshape(k, shape=(-1, 0, 0), reverse=True)
+    q = mx.sym.contrib.div_sqrt_dim(q)
+    att_score = mx.sym.batch_dot(q, k, transpose_b=True)
+    att_score = att_score + sonde
+    v = mx.sym.reshape(v, shape=(0, 0, num_heads, -1))
+    v = mx.sym.transpose(v, axes=(0, 2, 1, 3))
+    v = mx.sym.reshape(v, shape=(-1, 0, 0), reverse=True)
+    weighted_value = mx.sym.batch_dot(att_score, v)
+    weighted_value = mx.sym.reshape(weighted_value, shape=(-1, num_heads, 0, 0),
+                                    reverse=True)
+    weighted_value = mx.sym.transpose(weighted_value, axes=(0, 2, 1, 3))
+    weighted_value = mx.sym.reshape(weighted_value, shape=(0, 0, -1))
+    output = mx.sym.FullyConnected(weighted_value, weight=out_weight, bias=out_bias, flatten=False,
+                                   num_hidden=out_dim, no_bias=False)
+    output = mx.sym.Group([output, att_score])
+    executor = output.simple_bind(ctx=mx.gpu(0),
+                                  q=(batch_size, qkv_length, qkv_dim),
+                                  kv=(batch_size, qkv_length, qkv_dim),
+                                  type_dict={'q': dtype,
+                                             'kv': dtype},
+                                  grad_req='write', force_rebind=True)
+    executor.copy_params_from(arg_params, {})
+    executor.arg_dict['sonde'][:] = 0.
+    executor.arg_dict['sonde'].wait_to_read()
+    executor.forward(is_train=True)
+    output_orig = executor.outputs[0].asnumpy()
+    att_score_orig = executor.outputs[1].asnumpy()
+    executor.backward([mx.nd.array(output_grads, dtype=dtype), mx.nd.zeros(att_score_orig.shape, dtype=dtype)])
+    grads_orig = {k : v.asnumpy() for k, v in executor.grad_dict.items()}
+    assert_allclose(att_score_orig, att_score_opti, rtol=1e-2, atol=1e-3)
+    assert_allclose(output_orig, output_opti, rtol=1e-2, atol=1e-3)
+
+    for k in grads_opti.keys():
+        assert(grads_orig[k].dtype == grads_opti[k].dtype)
+        assert(grads_orig[k].shape == grads_opti[k].shape)
+        assert_allclose(grads_orig[k], grads_opti[k], rtol=1e-2, atol=1e-3)
+
+def test_multihead_attention_encdec():
+    for dtype in ['float16', 'float32']:
+        check_multihead_attention_encdec(dtype=dtype)
 
 if __name__ == '__main__':
     import nose
