[TE/JAX] XLA FFI calls for three cast transpose functions (#1310)

* FFI for some transpose & activation functions
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Remove comments in transformer_engine/jax/csrc/extensions/activation.cpp
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>
Signed-off-by: Hua Huang <huangh1994@outlook.com>

---------
Signed-off-by: Hua Huang <huah@nvidia.com>
Signed-off-by: Hua Huang <huangh1994@outlook.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Phuong Nguyen <36155692+phu0ngng@users.noreply.github.com>

tests/jax/test_custom_call_compute.py

@@ -22,6 +22,7 @@ from transformer_engine.jax.cpp_extensions.activation import _jax_act_lu
 from transformer_engine.jax.cpp_extensions.transpose import (
     _jax_transpose,
     _jax_cast_transpose,
+    _jax_dbias_cast_transpose,
 )
 from transformer_engine.jax.cpp_extensions.quantization import _jax_cast_fp8
 from transformer_engine.jax import cpp_extensions as tex
@@ -504,7 +505,6 @@ class TestActivationLuFP8(TestActivationLu):
                         scale_inv,
                         FP8Helper.BWD_DTYPE,
                         -1,
-                        -2,
                         self.activation_type,
                     )
                 )
@@ -812,6 +812,34 @@ class TestTranspose:
         assert_tree_like_allclose(jax_output, ffi_output)
         assert_tree_like_allclose(noffi_output, ffi_output)
 
+    @pytest.mark.parametrize(
+        "out_dtype",
+        [
+            pytest.param(jnp.float8_e4m3fn, id="output_float8_e4m3fn"),
+            pytest.param(jnp.float8_e5m2, id="output_float8_e5m2"),
+        ],
+    )
+    def test_dbias_cast_transpose(self, in_dtype, input_shape, transpose_axis, out_dtype):
+        amax = jnp.zeros(1, jnp.float32)
+        scale = jnp.ones(1, jnp.float32)
+        scale_inv = jnp.ones(1, jnp.float32)
+        key = jax.random.PRNGKey(0)
+        input = jax.random.uniform(key, input_shape, in_dtype)
+        static_axis_boundary = -1
+        jax_output = _jax_dbias_cast_transpose(
+            input, amax, scale, out_dtype, static_axis_boundary, transpose_axis
+        )
+        os.environ["NVTE_JAX_WITH_FFI"] = "0"
+        noffi_output = tex.dbias_cast_transpose(
+            input, amax, scale, scale_inv, out_dtype, static_axis_boundary, transpose_axis
+        )
+        os.environ["NVTE_JAX_WITH_FFI"] = "1"
+        ffi_output = tex.dbias_cast_transpose(
+            input, amax, scale, scale_inv, out_dtype, static_axis_boundary, transpose_axis
+        )
+        assert_tree_like_allclose(jax_output, ffi_output)
+        assert_tree_like_allclose(noffi_output, ffi_output)
+
 
 @pytest.mark.skipif(not is_fp8_supported, reason=reason)
 @pytest.mark.parametrize(

transformer_engine/jax/csrc/extensions.h

@@ -155,27 +155,29 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(TransposeHandler);
 
 void CastTranspose(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
 
+XLA_FFI_DECLARE_HANDLER_SYMBOL(CastTransposeHandler);
+
 pybind11::tuple GetDBiasCastTransposeWorkspaceSizes(size_t batch_size, size_t hidden_size,
                                                     DType in_dtype, DType out_dtype);
 
-XLA_FFI_DECLARE_HANDLER_SYMBOL(CastTransposeHandler);
-
 void DBiasCastTranspose(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
 
+XLA_FFI_DECLARE_HANDLER_SYMBOL(DBiasCastTransposeHandler);
+
 // Activation
 
 size_t get_activation_len(NVTE_Activation_Type activation_enum);
 
 void ActLu(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
 
-void ActLuFP8(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
-
-void DActLu(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
-
 XLA_FFI_DECLARE_HANDLER_SYMBOL(ActLuHandler);
 
+void ActLuFP8(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
+
 XLA_FFI_DECLARE_HANDLER_SYMBOL(ActLuFP8Handler);
 
+void DActLu(cudaStream_t stream, void **buffers, const char *opaque, size_t opaque_len);
+
 XLA_FFI_DECLARE_HANDLER_SYMBOL(DActLuHandler);
 
 pybind11::tuple GetDActDBiasCastTransposeWorkspaceSizes(size_t batch_size, size_t hidden_size,
@@ -184,9 +186,13 @@ pybind11::tuple GetDActDBiasCastTransposeWorkspaceSizes(size_t batch_size, size_
 void DActLuDBiasCastTranspose(cudaStream_t stream, void **buffers, const char *opaque,
                               size_t opaque_len);
 
+XLA_FFI_DECLARE_HANDLER_SYMBOL(DActLuDBiasCastTransposeHandler);
+
 void DGatedActLuCastTranspose(cudaStream_t stream, void **buffers, const char *opaque,
                               size_t opaque_len);
 
+XLA_FFI_DECLARE_HANDLER_SYMBOL(DGatedActLuCastTransposeHandler);
+
 // Normalization
 
 pybind11::tuple GetLayerNormForwardWorkspaceSizes(size_t batch_size, size_t hidden_size,

transformer_engine/jax/csrc/extensions/activation.cpp

@@ -373,7 +373,7 @@ void DActLuDBiasCastTranspose(cudaStream_t stream, void **buffers, const char *o
   auto m = desc.shape.dims[0];
   auto n = desc.shape.dims[1];
   auto act_enum = static_cast<NVTE_Activation_Type>(desc.act_enum);
-  ;
+
   auto input_shape = std::vector<size_t>{m, n};
   auto act_input_shape = std::vector<size_t>{m, n};
   auto output_shape = std::vector<size_t>{m, n};
@@ -422,6 +422,107 @@ void DActLuDBiasCastTranspose(cudaStream_t stream, void **buffers, const char *o
   }
 }
 
+Error_Type DActLuDBiasCastTransposeFFI(cudaStream_t stream, Buffer_Type input_buf,
+                                       Buffer_Type act_input_buf, Buffer_Type amax_buf,
+                                       Buffer_Type scale_buf, Buffer_Type scale_inv_buf,
+                                       Result_Type output_buf, Result_Type output_trans_buf,
+                                       Result_Type dbias_buf, Result_Type amax_out_buf,
+                                       Result_Type workspace_buf, int64_t act_enum) {
+  auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type());
+  auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type());
+  auto workspace_dtype = convert_ffi_datatype_to_te_dtype(workspace_buf->element_type());
+
+  auto *input = input_buf.untyped_data();
+  auto *act_input = act_input_buf.untyped_data();
+  float *amax = reinterpret_cast<float *>(amax_buf.untyped_data());
+  float *scale = reinterpret_cast<float *>(scale_buf.untyped_data());
+  float *scale_inv = reinterpret_cast<float *>(scale_inv_buf.untyped_data());
+  auto *output = output_buf->untyped_data();
+  auto *output_trans = output_trans_buf->untyped_data();
+  auto *dbias = dbias_buf->untyped_data();
+  float *amax_out = reinterpret_cast<float *>(amax_out_buf->untyped_data());
+  void *workspace = workspace_buf->untyped_data();
+  NVTE_CHECK(amax == amax_out,
+             "amax not bound to amax_out in TE/JAX DActLuDBiasCastTranspose primitive.");
+  if (!use_fp8(out_dtype)) {
+    scale = nullptr;
+    scale_inv = nullptr;
+    amax_out = nullptr;
+  }
+
+  auto input_dims = input_buf.dimensions();
+  auto act_input_dims = act_input_buf.dimensions();
+  auto workspace_dims = workspace_buf->dimensions();
+  // m = x_batch_size = reduce(operator.mul, x_shape[:-2]), x_shape == act_input_dims
+  // n = ir_dz_shape[-1], ir_dz_shape == input_dims
+  auto input_ranks = input_dims.size();
+  auto m = product(act_input_dims, 0, act_input_dims.size() - 2);
+  auto n = product(input_dims, input_ranks - 1, input_ranks);
+  auto input_shape = std::vector<size_t>{m, n};
+  auto act_input_shape = std::vector<size_t>{m, n};
+  auto output_shape = std::vector<size_t>{m, n};
+  auto output_trans_shape = std::vector<size_t>{n, m};
+  auto dbias_shape = std::vector<size_t>{n};
+  std::vector<size_t> workspace_shape(workspace_dims.begin(), workspace_dims.end());
+
+  auto input_tensor = TensorWrapper(input, input_shape, in_dtype);
+  auto act_input_tensor = TensorWrapper(act_input, input_shape, in_dtype);
+  auto output_tensor = TensorWrapper(output, output_shape, out_dtype, amax_out, scale, scale_inv);
+  auto output_trans_tensor =
+      TensorWrapper(output_trans, output_trans_shape, out_dtype, amax_out, scale, scale_inv);
+  auto dbias_tensor = TensorWrapper(dbias, dbias_shape, in_dtype);
+  auto workspace_tensor = TensorWrapper(workspace, workspace_shape, workspace_dtype);
+
+  auto act_type = static_cast<NVTE_Activation_Type>(act_enum);
+  switch (act_type) {
+    case NVTE_Activation_Type::GELU:
+      nvte_cast_transpose_dbias_dgelu(input_tensor.data(), act_input_tensor.data(),
+                                      output_tensor.data(), output_trans_tensor.data(),
+                                      dbias_tensor.data(), workspace_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::SILU:
+      nvte_cast_transpose_dbias_dsilu(input_tensor.data(), act_input_tensor.data(),
+                                      output_tensor.data(), output_trans_tensor.data(),
+                                      dbias_tensor.data(), workspace_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::RELU:
+      nvte_cast_transpose_dbias_drelu(input_tensor.data(), act_input_tensor.data(),
+                                      output_tensor.data(), output_trans_tensor.data(),
+                                      dbias_tensor.data(), workspace_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::QGELU:
+      nvte_cast_transpose_dbias_dqgelu(input_tensor.data(), act_input_tensor.data(),
+                                       output_tensor.data(), output_trans_tensor.data(),
+                                       dbias_tensor.data(), workspace_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::SRELU:
+      nvte_cast_transpose_dbias_dsrelu(input_tensor.data(), act_input_tensor.data(),
+                                       output_tensor.data(), output_trans_tensor.data(),
+                                       dbias_tensor.data(), workspace_tensor.data(), stream);
+      break;
+    default:
+      NVTE_ERROR("Unsupported ActivationEnum");
+      break;
+  }
+  return ffi_with_cuda_error_check();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(DActLuDBiasCastTransposeHandler, DActLuDBiasCastTransposeFFI,
+                              FFI::Bind()
+                                  .Ctx<FFI_Stream_Type>()  // stream
+                                  .Arg<Buffer_Type>()      // input
+                                  .Arg<Buffer_Type>()      // act_input
+                                  .Arg<Buffer_Type>()      // amax
+                                  .Arg<Buffer_Type>()      // scale
+                                  .Arg<Buffer_Type>()      // scale_inv
+                                  .Ret<Buffer_Type>()      // output
+                                  .Ret<Buffer_Type>()      // output_trans
+                                  .Ret<Buffer_Type>()      // dbias
+                                  .Ret<Buffer_Type>()      // amax_out
+                                  .Ret<Buffer_Type>()      // workspace
+                                  .Attr<int64_t>("act_enum"),
+                              FFI_CudaGraph_Traits);
+
 void DGatedActLuCastTranspose(cudaStream_t stream, void **buffers, const char *opaque,
                               size_t opaque_len) {
   auto *input = buffers[0];
@@ -444,7 +545,7 @@ void DGatedActLuCastTranspose(cudaStream_t stream, void **buffers, const char *o
   auto m = desc.shape.dims[0];
   auto n = desc.shape.dims[1];
   auto act_enum = static_cast<NVTE_Activation_Type>(desc.act_enum);
-  ;
+
   auto input_shape = desc.shape.to_vector();
   auto act_input_shape = std::vector<size_t>{m, n * 2};
   auto output_shape = std::vector<size_t>{m, n * 2};
@@ -484,5 +585,88 @@ void DGatedActLuCastTranspose(cudaStream_t stream, void **buffers, const char *o
   }
 }
 
+Error_Type DGatedActLuCastTransposeFFI(cudaStream_t stream, Buffer_Type input_buf,
+                                       Buffer_Type act_input_buf, Buffer_Type amax_buf,
+                                       Buffer_Type scale_buf, Buffer_Type scale_inv_buf,
+                                       Result_Type output_buf, Result_Type output_trans_buf,
+                                       Result_Type amax_out_buf, int64_t act_enum) {
+  auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type());
+  auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type());
+
+  auto *input = input_buf.untyped_data();
+  auto *act_input = act_input_buf.untyped_data();
+  float *amax = reinterpret_cast<float *>(amax_buf.untyped_data());
+  float *scale = reinterpret_cast<float *>(scale_buf.untyped_data());
+  float *scale_inv = reinterpret_cast<float *>(scale_inv_buf.untyped_data());
+  auto *output = output_buf->untyped_data();
+  auto *output_trans = output_trans_buf->untyped_data();
+  float *amax_out = reinterpret_cast<float *>(amax_out_buf->untyped_data());
+  NVTE_CHECK(amax == amax_out,
+             "amax not bound to amax_out in TE/JAX DGatedActLuCastTranspose primitive.");
+  if (!use_fp8(out_dtype)) {
+    scale = nullptr;
+    scale_inv = nullptr;
+    amax_out = nullptr;
+  }
+
+  auto input_dims = input_buf.dimensions();
+  auto act_input_dims = act_input_buf.dimensions();
+  auto act_input_ranks = act_input_dims.size();
+  auto m = product(act_input_dims, 0, act_input_ranks - 2);
+  auto n = product(act_input_dims, act_input_ranks - 1, act_input_ranks);
+  auto input_shape = std::vector<size_t>{m, n};
+  auto act_input_shape = std::vector<size_t>{m, n * 2};
+  auto output_shape = std::vector<size_t>{m, n * 2};
+  auto output_trans_shape = std::vector<size_t>{n * 2, m};
+
+  auto input_tensor = TensorWrapper(input, input_shape, in_dtype);
+  auto act_input_tensor = TensorWrapper(act_input, act_input_shape, in_dtype);
+  auto output_tensor = TensorWrapper(output, output_shape, out_dtype, amax_out, scale, scale_inv);
+  auto output_trans_tensor =
+      TensorWrapper(output_trans, output_trans_shape, out_dtype, amax_out, scale, scale_inv);
+
+  auto act_type = static_cast<NVTE_Activation_Type>(act_enum);
+  switch (act_type) {
+    case NVTE_Activation_Type::GEGLU:
+      nvte_dgeglu_cast_transpose(input_tensor.data(), act_input_tensor.data(), output_tensor.data(),
+                                 output_trans_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::SWIGLU:
+      nvte_dswiglu_cast_transpose(input_tensor.data(), act_input_tensor.data(),
+                                  output_tensor.data(), output_trans_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::REGLU:
+      nvte_dreglu_cast_transpose(input_tensor.data(), act_input_tensor.data(), output_tensor.data(),
+                                 output_trans_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::QGEGLU:
+      nvte_dqgeglu_cast_transpose(input_tensor.data(), act_input_tensor.data(),
+                                  output_tensor.data(), output_trans_tensor.data(), stream);
+      break;
+    case NVTE_Activation_Type::SREGLU:
+      nvte_dsreglu_cast_transpose(input_tensor.data(), act_input_tensor.data(),
+                                  output_tensor.data(), output_trans_tensor.data(), stream);
+      break;
+    default:
+      NVTE_ERROR("Unsupported ActivationEnum");
+      break;
+  }
+  return ffi_with_cuda_error_check();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(DGatedActLuCastTransposeHandler, DGatedActLuCastTransposeFFI,
+                              FFI::Bind()
+                                  .Ctx<FFI_Stream_Type>()  // stream
+                                  .Arg<Buffer_Type>()      // input
+                                  .Arg<Buffer_Type>()      // act_input
+                                  .Arg<Buffer_Type>()      // amax
+                                  .Arg<Buffer_Type>()      // scale
+                                  .Arg<Buffer_Type>()      // scale_inv
+                                  .Ret<Buffer_Type>()      // output
+                                  .Ret<Buffer_Type>()      // output_trans
+                                  .Ret<Buffer_Type>()      // amax_out
+                                  .Attr<int64_t>("act_enum"),
+                              FFI_CudaGraph_Traits);
+
 }  // namespace jax
 }  // namespace transformer_engine

transformer_engine/jax/csrc/extensions/pybind.cpp

@@ -55,11 +55,16 @@ pybind11::dict Registrations() {
   // Transpose
   dict["te_transpose_ffi"] = EncapsulateFFI(TransposeHandler);
   dict["te_cast_transpose_ffi"] = EncapsulateFFI(CastTransposeHandler);
+  dict["te_dbias_cast_transpose_ffi"] = EncapsulateFFI(DBiasCastTransposeHandler);
 
   // Activation
   dict["te_act_lu_ffi"] = EncapsulateFFI(ActLuHandler);
   dict["te_act_lu_fp8_ffi"] = EncapsulateFFI(ActLuFP8Handler);
   dict["te_dact_lu_ffi"] = EncapsulateFFI(DActLuHandler);
+  dict["te_dact_lu_dbias_cast_transpose_ffi"] =
+      EncapsulateFunction(DActLuDBiasCastTransposeHandler);
+  dict["te_dgated_act_lu_cast_transpose_ffi"] =
+      EncapsulateFunction(DGatedActLuCastTransposeHandler);
 
   // Quantization
   dict["te_quantize_ffi"] = EncapsulateFFI(QuantizeHandler);

transformer_engine/jax/csrc/extensions/transpose.cpp

@@ -100,18 +100,18 @@ void CastTranspose(cudaStream_t stream, void **buffers, const char *opaque, size
 
 Error_Type CastTransposeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type amax_buf,
                             Buffer_Type scale_buf, Buffer_Type scale_inv_buf,
-                            Result_Type input_cast_buf, Result_Type input_cast_trans_buf,
+                            Result_Type output_buf, Result_Type output_trans_buf,
                             Result_Type amax_out_buf, int64_t transpose_axis) {
   auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type());
-  auto out_dtype = convert_ffi_datatype_to_te_dtype(input_cast_buf->element_type());
+  auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type());
 
   auto *input = input_buf.untyped_data();
   float *amax = reinterpret_cast<float *>(amax_buf.untyped_data());
   float *scale = reinterpret_cast<float *>(scale_buf.untyped_data());
   float *scale_inv = reinterpret_cast<float *>(scale_inv_buf.untyped_data());
 
-  auto *input_cast = input_cast_buf->untyped_data();
-  auto *input_cast_trans = input_cast_trans_buf->untyped_data();
+  auto *output = output_buf->untyped_data();
+  auto *output_trans = output_trans_buf->untyped_data();
   float *amax_out = reinterpret_cast<float *>(amax_out_buf->untyped_data());
   NVTE_CHECK(amax == amax_out, "amax not bound to amax_out in TE/JAX CastTranspose primitive.");
 
@@ -126,15 +126,15 @@ Error_Type CastTransposeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_T
   auto m = product(input_dims, 0, transpose_axis);
   auto n = product(input_dims, transpose_axis, input_dims.size());
   auto input_shape = std::vector<size_t>{m, n};
-  auto input_trans_shape = std::vector<size_t>{n, m};
+  auto output_shape = input_shape;
+  auto output_trans_shape = std::vector<size_t>{n, m};
 
   auto input_tensor = TensorWrapper(input, input_shape, in_dtype);
-  auto input_cast_tensor =
-      TensorWrapper(input_cast, input_shape, out_dtype, amax_out, scale, scale_inv);
-  auto input_cast_trans_tensor =
-      TensorWrapper(input_cast_trans, input_trans_shape, out_dtype, amax_out, scale, scale_inv);
+  auto output_tensor = TensorWrapper(output, output_shape, out_dtype, amax_out, scale, scale_inv);
+  auto output_trans_tensor =
+      TensorWrapper(output_trans, output_trans_shape, out_dtype, amax_out, scale, scale_inv);
 
-  nvte_cast_transpose(input_tensor.data(), input_cast_tensor.data(), input_cast_trans_tensor.data(),
+  nvte_cast_transpose(input_tensor.data(), output_tensor.data(), output_trans_tensor.data(),
                       stream);
   return ffi_with_cuda_error_check();
 }
@@ -146,8 +146,8 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(CastTransposeHandler, CastTransposeFFI,
                                   .Arg<Buffer_Type>()      // amax
                                   .Arg<Buffer_Type>()      // scale
                                   .Arg<Buffer_Type>()      // scale_inv
-                                  .Ret<Buffer_Type>()      // input_cast
-                                  .Ret<Buffer_Type>()      // input_cast_trans
+                                  .Ret<Buffer_Type>()      // output
+                                  .Ret<Buffer_Type>()      // output_trans
                                   .Ret<Buffer_Type>()      // amax_out
                                   .Attr<int64_t>("transpose_axis"),
                               FFI_CudaGraph_Traits);
@@ -213,5 +213,70 @@ void DBiasCastTranspose(cudaStream_t stream, void **buffers, const char *opaque,
                             dbias_tensor.data(), workspace.data(), stream);
 }
 
+Error_Type DBiasCastTransposeFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type amax_buf,
+                                 Buffer_Type scale_buf, Buffer_Type scale_inv_buf,
+                                 Result_Type output_buf, Result_Type output_trans_buf,
+                                 Result_Type dbias_buf, Result_Type amax_out_buf,
+                                 Result_Type workspace_buf, int64_t transpose_axis) {
+  auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type());
+  auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type());
+  auto workspace_dtype = convert_ffi_datatype_to_te_dtype(workspace_buf->element_type());
+
+  auto *input = input_buf.untyped_data();
+  float *amax = reinterpret_cast<float *>(amax_buf.untyped_data());
+  float *scale = reinterpret_cast<float *>(scale_buf.untyped_data());
+  float *scale_inv = reinterpret_cast<float *>(scale_inv_buf.untyped_data());
+
+  auto *output = output_buf->untyped_data();
+  auto *output_trans = output_trans_buf->untyped_data();
+  auto *dbias = dbias_buf->untyped_data();
+  float *amax_out = reinterpret_cast<float *>(amax_out_buf->untyped_data());
+  void *workspace = workspace_buf->untyped_data();
+  NVTE_CHECK(amax == amax_out,
+             "amax not bound to amax_out in TE/JAX DBiasCastTranspose primitive.");
+  if (!use_fp8(out_dtype)) {
+    scale = nullptr;
+    scale_inv = nullptr;
+    amax_out = nullptr;
+  }
+
+  auto input_dims = input_buf.dimensions();
+  auto workspace_dims = workspace_buf->dimensions();
+  if (transpose_axis < 0) transpose_axis += input_dims.size();
+  auto m = product(input_dims, 0, transpose_axis);
+  auto n = product(input_dims, transpose_axis, input_dims.size());
+  auto input_shape = std::vector<size_t>{m, n};
+  auto output_shape = std::vector<size_t>{m, n};
+  auto output_trans_shape = std::vector<size_t>{n, m};
+  auto dbias_shape = std::vector<size_t>{n};
+  std::vector<size_t> workspace_shape(workspace_dims.begin(), workspace_dims.end());
+
+  auto input_tensor = TensorWrapper(input, input_shape, in_dtype);
+  auto output_tensor = TensorWrapper(output, output_shape, out_dtype, amax_out, scale, scale_inv);
+  auto output_trans_tensor =
+      TensorWrapper(output_trans, output_trans_shape, out_dtype, amax_out, scale, scale_inv);
+  auto dbias_tensor = TensorWrapper(dbias, dbias_shape, in_dtype);
+  auto workspace_tensor = TensorWrapper(workspace, workspace_shape, workspace_dtype);
+
+  nvte_cast_transpose_dbias(input_tensor.data(), output_tensor.data(), output_trans_tensor.data(),
+                            dbias_tensor.data(), workspace_tensor.data(), stream);
+  return ffi_with_cuda_error_check();
+}
+
+XLA_FFI_DEFINE_HANDLER_SYMBOL(DBiasCastTransposeHandler, DBiasCastTransposeFFI,
+                              FFI::Bind()
+                                  .Ctx<FFI_Stream_Type>()  // stream
+                                  .Arg<Buffer_Type>()      // input
+                                  .Arg<Buffer_Type>()      // amax
+                                  .Arg<Buffer_Type>()      // scale
+                                  .Arg<Buffer_Type>()      // scale_inv
+                                  .Ret<Buffer_Type>()      // output
+                                  .Ret<Buffer_Type>()      // output_trans
+                                  .Ret<Buffer_Type>()      // dbias
+                                  .Ret<Buffer_Type>()      // amax_out
+                                  .Ret<Buffer_Type>()      // workspace
+                                  .Attr<int64_t>("transpose_axis"),
+                              FFI_CudaGraph_Traits);
+
 }  // namespace jax
 }  // namespace transformer_engine

transformer_engine/jax/layernorm_mlp.py

@@ -516,7 +516,6 @@ def _fused_layernorm_fp8_mlp_bwd_rule(
                     dactivation_lu_scale_inv,
                     bwd_dtype,
                     static_axis_boundary=-1,
-                    transpose_axis_boundary=-2,
                     activation_type=activation_type,
                 )
             )