unify with nvte_multi_tensor_quantize

Author: alextmagro

benchmarks/cpp/cast/bench_group_quantize_mxfp8.cpp

@@ -303,7 +303,7 @@ static void BM_MultiQuantizeMXFP8(benchmark::State &state) {
 
   for (auto _ : state) {
     HIP_CHECK(hipEventRecord(start, stream));
-    nvte_multi_quantize_mxfp8(num_experts, nvte_inputs.data(), nvte_outputs.data(), stream);
+    nvte_multi_tensor_quantize(nvte_inputs.data(), nvte_outputs.data(), nullptr, num_experts, stream);
     HIP_CHECK(hipEventRecord(stop, stream));
     HIP_CHECK(hipEventSynchronize(stop));
     float ms = 0;

tests/cpp/operator/test_multi_quantize_mxfp8.cu

@@ -47,8 +47,8 @@ void performTest(const std::vector<std::pair<size_t, size_t>> &tensor_dims,
     nvte_outputs_multi.push_back(t.data());
   }
 
-  nvte_multi_quantize_mxfp8(num_tensors, nvte_inputs.data(),
-                            nvte_outputs_multi.data(), 0);
+  nvte_multi_tensor_quantize(nvte_inputs.data(), nvte_outputs_multi.data(),
+                             nullptr, num_tensors, 0);
 
   for (size_t i = 0; i < num_tensors; i++) {
     if (tensor_dims[i].first > 0 && tensor_dims[i].second > 0)

transformer_engine/common/cast/cast.cu

@@ -30,20 +30,6 @@ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t strea
   dispatch::quantize_fwd_helper<IS_ACT, Empty, nullptr>(input, output, nullptr, stream);
 }
 
-#ifdef __HIP_PLATFORM_AMD__
-void nvte_multi_quantize_mxfp8(size_t num_tensors, const NVTETensor *input_list,
-                               NVTETensor *output_list, cudaStream_t stream) {
-  NVTE_API_CALL(nvte_multi_quantize_mxfp8);
-  using namespace transformer_engine;
-  std::vector<Tensor *> input_list_, output_list_;
-  for (size_t i = 0; i < num_tensors; i++) {
-    input_list_.push_back(convertNVTETensorCheck(input_list[i]));
-    output_list_.push_back(convertNVTETensorCheck(output_list[i]));
-  }
-  dispatch::multi_quantize_mxfp8(input_list_, output_list_, stream);
-}
-#endif
-
 void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output,
                          cudaStream_t stream) {
   NVTE_API_CALL(nvte_group_quantize);
@@ -113,6 +99,19 @@ void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs,
   NVTE_API_CALL(nvte_multi_tensor_quantize);
   using namespace transformer_engine;
 
+#ifdef __HIP_PLATFORM_AMD__
+  if (num_tensors > 0 &&
+      convertNVTETensorCheck(outputs[0])->scaling_mode == NVTE_MXFP8_1D_SCALING) {
+    std::vector<Tensor *> input_list, output_list;
+    for (size_t i = 0; i < num_tensors; i++) {
+      input_list.push_back(convertNVTETensorCheck(inputs[i]));
+      output_list.push_back(convertNVTETensorCheck(outputs[i]));
+    }
+    dispatch::multi_quantize_mxfp8(input_list, output_list, stream);
+    return;
+  }
+#endif
+
   constexpr bool IS_ACT = false;
 
   const size_t num_streams = nvte_get_num_compute_streams();

transformer_engine/common/include/transformer_engine/cast.h

@@ -101,19 +101,6 @@ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t strea
 void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output,
                          cudaStream_t stream);
 
-#ifdef __HIP_PLATFORM_AMD__
-/*! \brief Fused multi-tensor MXFP8 quantize. Quantizes multiple tensors in a single kernel launch.
- *         Each tensor can have different shapes. Output tensors are written to per-tensor pointers.
- *
- *  \param[in]     num_tensors      Number of tensors to quantize.
- *  \param[in]     input_list       Array of input tensors.
- *  \param[in,out] output_list      Array of output MXFP8 tensors.
- *  \param[in]     stream           CUDA stream used for the operation.
- */
-void nvte_multi_quantize_mxfp8(size_t num_tensors, const NVTETensor *input_list,
-                               NVTETensor *output_list, cudaStream_t stream);
-#endif
-
 /*! \brief Casts input tensor to FP8/MXFP8/BlockwiseFP8, providing the option to immediately exit the kernel
  *         based on the value of the 'noop' tensor.
  *         The type of quantized tensor in the output depends on the scaling mode of the output

transformer_engine/pytorch/csrc/extensions/cast.cpp

@@ -291,20 +291,6 @@ void multi_tensor_quantize_impl(const std::vector<TensorWrapper> &input_list,
   }
 
   // Launch TE kernel
-#ifdef USE_ROCM
-  if (num_tensors > 0 && detail::IsMXFP8Quantizers(quantizer_py_list[0].ptr())) {
-    std::vector<NVTETensor> nvte_input_list, nvte_output_list;
-    for (size_t i = 0; i < num_tensors; i++) {
-      nvte_input_list.push_back(input_list[i].data());
-      nvte_output_list.push_back(output_list[i].data());
-    }
-    NVTE_SCOPED_GIL_RELEASE({
-      nvte_multi_quantize_mxfp8(nvte_input_list.size(), nvte_input_list.data(),
-                                nvte_output_list.data(), at::cuda::getCurrentCUDAStream());
-    });
-    return;
-  }
-#endif
   if (with_fused_kernel) {
     // Fused kernel for multi-tensor quantize
     std::vector<NVTETensor> nvte_tensor_input_list;
@@ -317,6 +303,19 @@ void multi_tensor_quantize_impl(const std::vector<TensorWrapper> &input_list,
       nvte_multi_cast_transpose(nvte_tensor_input_list.size(), nvte_tensor_input_list.data(),
                                 nvte_tensor_output_list.data(), at::cuda::getCurrentCUDAStream());
     });
+#ifdef USE_ROCM
+  } else if (num_tensors > 0 && detail::IsMXFP8Quantizers(quantizer_py_list[0].ptr())) {
+    std::vector<NVTETensor> nvte_tensor_input_list;
+    std::vector<NVTETensor> nvte_tensor_output_list;
+    for (size_t i = 0; i < num_tensors; ++i) {
+      nvte_tensor_input_list.push_back(input_list[i].data());
+      nvte_tensor_output_list.push_back(output_list[i].data());
+    }
+    NVTE_SCOPED_GIL_RELEASE({
+      nvte_multi_tensor_quantize(nvte_tensor_input_list.data(), nvte_tensor_output_list.data(),
+                                 nullptr, num_tensors, at::cuda::getCurrentCUDAStream());
+    });
+#endif
   } else {
     // Quantize kernels individually
     for (size_t i = 0; i < num_tensors; ++i) {