Enable MultiCastTranspose for expert weights

transformer_engine/pytorch/csrc/extensions.h

@@ -317,7 +317,13 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 const size_t num_tensors, std::optional<at::Tensor> first_dims);
 
 std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
-                                              std::vector<py::handle> quantizer_list);
+                                              std::vector<py::handle> quantizer_list,
+                                              // ROCm adds a fused multi-quantize
+                                              // path for GroupedLinear expert weights. The cached
+                                              // workspace can be passed directly as `outputs`, so the
+                                              // kernel writes in place instead of allocating separate
+                                              // storage.
+                                              const py::object &outputs = py::none());
 
 std::vector<py::object> split_quantize(const at::Tensor &tensor,
                                        const std::vector<size_t> &split_sections,

transformer_engine/pytorch/csrc/extensions/cast.cpp

@@ -491,12 +491,23 @@ void multi_tensor_quantize_impl(const std::vector<TensorWrapper> &input_list,
 }  // namespace
 
 std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
-                                              std::vector<py::handle> quantizer_list) {
+                                              std::vector<py::handle> quantizer_list,
+                                              const py::object &outputs) {
   // Check number of tensors
   const size_t num_tensors = tensor_list.size();
   NVTE_CHECK(quantizer_list.size() == num_tensors, "Expected ", num_tensors,
              " quantizers, but got ", quantizer_list.size());
 
+  const bool use_provided_outputs = !outputs.is_none();
+  py::sequence outputs_seq;
+  if (use_provided_outputs) {
+    NVTE_CHECK(py::isinstance<py::list>(outputs) || py::isinstance<py::tuple>(outputs),
+               "multi_tensor_quantize: outputs must be None, a list, or a tuple.");
+    outputs_seq = py::reinterpret_borrow<py::sequence>(outputs);
+    NVTE_CHECK(static_cast<size_t>(outputs_seq.size()) == num_tensors, "multi_tensor_quantize: ",
+               "len(outputs) is ", outputs_seq.size(), " but expected ", num_tensors, ".");
+  }
+
   // Convert quantizers to C++ objects
   std::vector<std::unique_ptr<Quantizer>> quantizer_cpp_list;
   for (size_t i = 0; i < num_tensors; i++) {
@@ -516,9 +527,17 @@ std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &ten
     const auto input_shape = input_cpp.shape();
     const auto input_dtype = GetTransformerEngineDType(input_py.scalar_type());
 
-    // Construct output tensor
-    std::vector<size_t> output_shape(input_shape.data, input_shape.data + input_shape.ndim);
-    auto [output_cpp, output_py] = quantizer_cpp_list[i]->create_tensor(output_shape, input_dtype);
+    TensorWrapper output_cpp;
+    py::object output_py;
+    if (use_provided_outputs) {
+      py::object output_obj = outputs_seq[static_cast<ssize_t>(i)];
+      std::tie(output_cpp, output_py) =
+          quantizer_cpp_list[i]->convert_and_update_tensor(std::move(output_obj));
+    } else {
+      std::vector<size_t> output_shape(input_shape.data, input_shape.data + input_shape.ndim);
+      std::tie(output_cpp, output_py) =
+          quantizer_cpp_list[i]->create_tensor(output_shape, input_dtype);
+    }
     output_cpp_list.emplace_back(std::move(output_cpp));
     output_py_list.emplace_back(std::move(output_py));
   }

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -301,7 +301,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("rmsnorm_bwd_add", &transformer_engine::pytorch::rmsnorm_bwd_add,
         "Fused backward of RMSNorm + add");
   m.def("multi_tensor_quantize", &transformer_engine::pytorch::multi_tensor_quantize,
-        "Multi-tensor quantize", py::arg("tensor_list"), py::arg("quantizer_list"));
+        "Multi-tensor quantize", py::arg("tensor_list"), py::arg("quantizer_list"),
+        py::arg("outputs") = py::none());
   m.def("split_quantize", &transformer_engine::pytorch::split_quantize,
         "Split and multi-tensor quantize", py::arg("tensor"), py::arg("split_sections"),
         py::arg("quantizer_list"), py::arg("disable_bulk_allocation") = false

transformer_engine/pytorch/module/base.py

@@ -801,6 +801,127 @@ def quantize_weight(
     return out, None
 
 
+def quantize_multi_weight(
+    *,
+    tensors: List[torch.Tensor],
+    quantizers: List[Quantizer],
+    workspaces: Optional[List[Optional[QuantizedTensorStorage]]] = None,
+    update_workspace: bool = True,
+    skip_update_flag: Optional[torch.Tensor] = None,
+    workspace_dtype: Optional[torch.dtype] = None,
+    cache: bool = False,
+) -> Tuple[List[QuantizedTensorStorage], List[Optional[QuantizedTensorStorage]]]:
+    """Quantize a group of weights, optionally reusing cached workspaces.
+
+    Group analogue of :func:`quantize_weight`. For delayed-scaling FP8 the whole
+    group is cast and transposed with a single fused ``multi_cast_transpose``
+    kernel, and on ROCm an MXFP8 group is quantized with a single fused
+    multi-quantize kernel, instead of one quantize call per tensor. When fusion
+    is not applicable (other recipes, FP8 rowwise-only usage, already-quantized
+    weights, or a CUDA-graph skip flag) the call falls back to
+    :func:`quantize_weight` for each tensor.
+
+    Parameters
+    ----------
+    tensors: list of torch.Tensor
+        Weight tensors to quantize.
+    quantizers: list of Quantizer
+        Quantizers for casting the weights.
+    workspaces: list of QuantizedTensorStorage, optional
+        Previously cached workspaces (from the module's ``_fp8_workspaces``).
+        ``None`` entries indicate a cache miss.
+    update_workspace: bool, default = True
+        Whether to update existing workspaces with fresh values.
+    skip_update_flag: torch.Tensor, optional
+        GPU flag to conditionally skip the update.
+    workspace_dtype: torch.dtype, optional
+        High-precision dtype for debug quantization workspaces.
+    cache: bool, default = False
+        If ``True``, brand-new workspaces are returned so the caller can store
+        them.
+
+    Returns
+    -------
+    (weights, new_workspaces)
+        *weights*: quantized weights ready for GEMM.
+        *new_workspaces*: per-tensor entries that are non-``None`` only when a
+        workspace should be (re)stored in ``_fp8_workspaces`` (i.e. ``cache`` is
+        ``True`` and the fused/per-tensor path produced a workspace to cache).
+    """
+    num_tensors = len(tensors)
+    if workspaces is None:
+        workspaces = [None] * num_tensors
+
+    # Fused path eligibility. Two recipes have a fused multi-tensor quantize kernel:
+    #   * delayed-scaling FP8 (Float8Quantizer): cast + transpose, so it requires
+    #     columnwise usage.
+    #   * MXFP8 (ROCm only): fused multi-quantize kernel that handles whichever
+    #     rowwise/columnwise buffers are allocated.
+    # Both require high-precision (not already quantized) weights and no CUDA-graph
+    # skip flag (the fused kernels have no device-side noop and would not preserve
+    # cached buffer pointers).
+    fused_fp8 = all(
+        isinstance(quantizer, Float8Quantizer) and quantizer.columnwise_usage
+        for quantizer in quantizers
+    )
+    fused_mxfp8 = IS_HIP_EXTENSION and all(
+        isinstance(quantizer, MXFP8Quantizer) for quantizer in quantizers
+    )
+    can_fuse = (
+        num_tensors > 0
+        and skip_update_flag is None
+        and (fused_fp8 or fused_mxfp8)
+        and not any(isinstance(tensor, QuantizedTensorStorage) for tensor in tensors)
+    )
+
+    if can_fuse:
+        # Validate cached workspaces; treat any invalid/missing entry as a cache miss.
+        cache_miss = False
+        for i, (workspace, quantizer) in enumerate(zip(workspaces, quantizers)):
+            if workspace is not None and not _is_weight_workspace_valid(workspace, quantizer):
+                workspaces[i] = None
+            if workspaces[i] is None:
+                cache_miss = True
+
+        new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * num_tensors
+        if not cache_miss and not update_workspace:
+            # All workspaces valid and no refresh requested.
+            return list(workspaces), new_workspaces
+
+        # Force internal=False so cached workspaces survive prepare_for_saving.
+        if cache:
+            saved_internal = [quantizer.internal for quantizer in quantizers]
+            for quantizer in quantizers:
+                quantizer.internal = False
+        if cache_miss:
+            # Single fused kernel allocates all outputs.
+            weights = tex.multi_tensor_quantize(list(tensors), quantizers)
+        else:
+            # In-place refresh of cached workspaces.
+            weights = tex.multi_tensor_quantize(list(tensors), quantizers, outputs=list(workspaces))
+        if cache:
+            for quantizer, internal in zip(quantizers, saved_internal):
+                quantizer.internal = internal
+            new_workspaces = list(weights)
+        return list(weights), new_workspaces
+
+    # Fallback: quantize each weight individually.
+    weights = []
+    new_workspaces = [None] * num_tensors
+    for i in range(num_tensors):
+        weight, new_workspaces[i] = quantize_weight(
+            tensor=tensors[i],
+            quantizer=quantizers[i],
+            workspace=workspaces[i],
+            update_workspace=update_workspace,
+            skip_update_flag=skip_update_flag,
+            workspace_dtype=workspace_dtype,
+            cache=cache,
+        )
+        weights.append(weight)
+    return weights, new_workspaces
+
+
 class TransformerEngineBaseModule(torch.nn.Module, ABC):
     """Base TE module."""
 

transformer_engine/pytorch/module/grouped_linear.py

@@ -20,6 +20,7 @@
 from .base import (
     get_dummy_wgrad,
     quantize_weight,
+    quantize_multi_weight,
     TransformerEngineBaseModule,
     _2X_ACC_FPROP,
     _2X_ACC_DGRAD,
@@ -198,20 +199,20 @@ def forward(
         weights_fp8: list
         new_workspaces = [None] * num_gemms
         if fp8 or debug:
-            weights_fp8 = []
+            # FP8 cast to workspace buffer. For delayed-scaling FP8 the whole group is
+            # cast and transposed with a single fused multi_cast_transpose kernel, and on
+            # ROCm an MXFP8 group is quantized with a single fused multi-quantize kernel;
+            # other cases fall back to a per-weight quantize inside quantize_multi_weight.
             update_ws = is_first_microbatch is None or is_first_microbatch
-            for i in range(num_gemms):
-                weight_fp8, new_workspaces[i] = quantize_weight(
-                    tensor=weights[i],
-                    quantizer=weight_quantizers[i],
-                    workspace=weight_workspaces[i] if weight_workspaces else None,
-                    update_workspace=update_ws,
-                    skip_update_flag=skip_fp8_weight_update,
-                    workspace_dtype=activation_dtype,
-                    cache=cache_weight,
-                )
-                weights_fp8.append(weight_fp8)
-
+            weights_fp8, new_workspaces = quantize_multi_weight(
+                tensors=list(weights),
+                quantizers=weight_quantizers,
+                workspaces=weight_workspaces if weight_workspaces else [None] * num_gemms,
+                update_workspace=update_ws,
+                skip_update_flag=skip_fp8_weight_update,
+                workspace_dtype=activation_dtype,
+                cache=cache_weight,
+            )
         else:
             weights_fp8 = [cast_if_needed(weight, activation_dtype) for weight in weights]