[Pytorch] Add support for row-wise quanted input for grouped gemm

transformer_engine/pytorch/csrc/extensions.h

@@ -367,6 +367,11 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 std::optional<at::Tensor> last_dims,
                                 std::optional<at::Tensor> tensor_offsets);
 
+py::object group_requantize_columnwise_and_swizzle_rowwise_(
+    py::handle grouped_x, py::handle columnwise_quantizer, const size_t num_tensors,
+    std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
+    bool return_dequantized);
+
 std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
                                               std::vector<py::handle> quantizer_list);
 

transformer_engine/pytorch/csrc/extensions/cast.cpp

@@ -663,6 +663,66 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o
   return py::reinterpret_borrow<py::object>(out_py);
 }
 
+py::object group_requantize_columnwise_and_swizzle_rowwise_(
+    py::handle grouped_x, py::handle columnwise_quantizer, const size_t num_tensors,
+    std::optional<at::Tensor> first_dims, DType otype, std::optional<at::Tensor> tensor_offsets,
+    bool return_dequantized) {
+  init_extension();
+
+  NVTE_CHECK(!grouped_x.attr("rowwise_data").is_none(),
+             "Pre-quantized MXFP8 grouped input is missing rowwise data.");
+  NVTE_CHECK(!grouped_x.attr("scale_inv").is_none(),
+             "Pre-quantized MXFP8 grouped input is missing rowwise scales.");
+  NVTE_CHECK(!grouped_x.attr("_with_gemm_swizzled_scales").cast<bool>(),
+             "Pre-quantized MXFP8 grouped input must have unswizzled scales.");
+  NVTE_CHECK(grouped_x.attr("columnwise_data").is_none(),
+             "Pre-quantized MXFP8 grouped input must be rowwise-only.");
+  if (!grouped_x.attr("quantizer").is_none()) {
+    // The GEMM consumes the input's rowwise data verbatim while the wgrad GEMM consumes the
+    // columnwise copy built here, so a dtype mismatch would make the two directions disagree.
+    NVTE_CHECK(grouped_x.attr("quantizer").attr("dtype").cast<DType>() ==
+                   columnwise_quantizer.attr("dtype").cast<DType>(),
+               "Pre-quantized MXFP8 grouped input and the columnwise quantizer disagree on the "
+               "FP8 dtype.");
+  }
+
+  const auto logical_shape = grouped_x.attr("logical_shape").cast<py::tuple>();
+  const auto total_tokens = logical_shape[0].cast<size_t>();
+  const auto hidden_dim = logical_shape[1].cast<size_t>();
+  // Each group's token count must be a multiple of 128 too, so that every group's scales start
+  // on a swizzle-tile boundary. Those counts live on the device (host reads would break CUDA
+  // graph capture), so that half is the caller's contract rather than an assertion.
+  NVTE_CHECK(total_tokens % 128 == 0 && hidden_dim % 128 == 0,
+             "Pre-quantized MXFP8 grouped input requires dims that are multiples of 128, but got (",
+             total_tokens, ", ", hidden_dim, ").");
+
+  // Dequantize first: it reads the rowwise scales, which the swizzle below replaces.
+  auto dequantized_grouped = group_dequantize(grouped_x, otype);
+  auto dequantized =
+      dequantized_grouped.attr("rowwise_data")
+          .cast<at::Tensor>()
+          .view({static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
+
+  // Swizzle the rowwise scales before attaching any columnwise data: a rowwise-only swizzle
+  // resets columnwise_scale_inv to None, which would strand the columnwise data below with a
+  // null scale pointer.
+  grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false);
+
+  // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
+  // data because the two directions scale along perpendicular axes. The caller hands us a
+  // columnwise-only, optimize_for_gemm quantizer, so the kernel emits swizzled columnwise
+  // scales directly.
+  auto columnwise = group_quantize(dequantized, columnwise_quantizer, num_tensors, first_dims,
+                                   std::nullopt, tensor_offsets, std::nullopt);
+  grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
+  grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
+
+  if (return_dequantized) {
+    return py::cast(dequantized);
+  }
+  return py::none();
+}
+
 namespace {
 
 void multi_tensor_quantize_impl(const std::vector<TensorWrapper> &input_list,

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -216,6 +216,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
         py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none());
+  m.def("group_requantize_columnwise_and_swizzle_rowwise_",
+        transformer_engine::pytorch::group_requantize_columnwise_and_swizzle_rowwise_,
+        "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle "
+        "its rowwise scales for GEMM, in place",
+        py::arg("grouped_x"), py::arg("columnwise_quantizer"), py::arg("num_tensors"),
+        py::arg("first_dims"), py::arg("otype"), py::arg("tensor_offsets") = py::none(),
+        py::arg("return_dequantized") = false);
   m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize,
         "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer"));
   m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)",

transformer_engine/pytorch/csrc/extensions/swizzle.cpp

@@ -389,6 +389,23 @@ std::optional<SwizzledGroupedScales> maybe_swizzle_grouped_tensor(GroupedTensorW
         tensor_offsets.data_ptr, static_cast<DType>(tensor_offsets.dtype), tensor_offsets.shape);
   }
 
+  // Varying per-tensor dimensions. Leaving these unset declares the grouped tensor uniform,
+  // which selects the uniform-shape swizzle kernel.
+  const auto first_dims = input.get_first_dims();
+  if (first_dims.data_ptr != nullptr) {
+    swizzle_input.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
+                                 first_dims.shape);
+    swizzle_output.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
+                                  first_dims.shape);
+  }
+  const auto last_dims = input.get_last_dims();
+  if (last_dims.data_ptr != nullptr) {
+    swizzle_input.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
+                                last_dims.shape);
+    swizzle_output.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
+                                 last_dims.shape);
+  }
+
   // Per-tensor logical dimensions (uniform-shape grouped tensor).
   const size_t num_tensors = input.num_tensors();
   const auto logical_shape_nvte = input.logical_shape();

transformer_engine/pytorch/ops/_common.py

@@ -14,6 +14,7 @@
 from ..torch_version import torch_version
 from ..quantization import FP8GlobalStateManager
 from ..tensor.float8_tensor import Float8Tensor
+from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..quantized_tensor import QuantizedTensorStorage
 from ..utils import canonicalize_dtype
 
@@ -66,6 +67,15 @@ def maybe_dequantize(
     return tensor
 
 
+def make_columnwise_gemm_quantizer(quantizer: MXFP8Quantizer) -> MXFP8Quantizer:
+    """Copy of ``quantizer`` configured to emit only GEMM-swizzled columnwise data."""
+    columnwise_quantizer = quantizer.copy()
+    columnwise_quantizer.set_usage(rowwise=False, columnwise=True)
+    columnwise_quantizer.optimize_for_gemm = True
+    columnwise_quantizer.internal = True
+    return columnwise_quantizer
+
+
 def maybe_autocast_dtype(
     *,
     device_type: str = "cuda",

transformer_engine/pytorch/ops/basic/grouped_linear.py

@@ -14,7 +14,7 @@
 import torch
 
 import transformer_engine_torch as tex
-from ...constants import DType
+from ...constants import DType, TE_DType
 from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor
 from ...distributed import CudaRNGStatesTracker
 from ...module._common import WeightGradStore
@@ -48,6 +48,7 @@
     get_dummy_wgrads_for_params,
     get_main_grad_from_param,
     is_quantized_tensor,
+    make_columnwise_gemm_quantizer,
     maybe_dequantize,
     validate_or_alloc_output,
     view_main_grad_as_grouped_buffer,
@@ -1201,6 +1202,11 @@ def _fuser_forward_split_quantize(
         out_buffer: Optional[torch.Tensor] = None,
     ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]:
         """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow."""
+        if isinstance(input_, GroupedTensor):
+            raise NotImplementedError(
+                "Pre-quantized GroupedTensor input is only supported on the "
+                "graph-safe grouped-tensor path."
+            )
         num_groups = self.num_groups
         has_bias = self.has_bias
 
@@ -1325,11 +1331,46 @@ def _fuser_forward_grouped_tensor(
 
         # Flatten to 2D so the first dim is the total token count.
         original_shape = list(input_.size())
-        x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features)
-        total_tokens = x.size(0)
+        prequantized_mxfp8_input = (
+            with_quantized_compute
+            and isinstance(input_, GroupedTensor)
+            and isinstance(input_quantizers[0], MXFP8Quantizer)
+            and isinstance(input_.quantizer, MXFP8Quantizer)
+        )
+        if prequantized_mxfp8_input:
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, in_features) layout; just validate the shape.
+            if input_.dim() != 2 or input_.size(-1) != self.in_features:
+                raise ValueError(
+                    "GroupedTensor input must have shape (total_tokens, "
+                    f"{self.in_features}), but got {tuple(input_.size())}."
+                )
+            total_tokens = input_.size(0)
+        else:
+            x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features)
+            total_tokens = x.size(0)
 
         # Build the input GroupedTensor.
-        if with_quantized_compute:
+        if prequantized_mxfp8_input:
+            # Rowwise-only MXFP8 input (e.g. FP8 token dispatch): feed the
+            # rowwise data to the forward GEMM as-is, manufacture the
+            # columnwise copy needed by the wgrad GEMM, and swizzle the
+            # rowwise scales for the GEMM.
+            grouped_x = input_.copy()
+            if weight_requires_grad:
+                tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                    grouped_x,
+                    make_columnwise_gemm_quantizer(input_quantizers[0]),
+                    num_groups,
+                    split_sizes,
+                    TE_DType[dtype],
+                    tensor_offsets=base_split_offsets * self.in_features,
+                )
+            else:
+                # No wgrad, so no columnwise copy is needed. The forward GEMM
+                # still requires swizzled rowwise scales.
+                tex.grouped_swizzle_for_gemm(grouped_x, rowwise=True, columnwise=False)
+        elif with_quantized_compute:
             input_quantizer = input_quantizers[0]
             input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
             input_quantizer.optimize_for_gemm = True
@@ -1682,8 +1723,25 @@ def _fuser_backward_grouped_tensor(
 
         # Flatten grad_output to 2D (total_tokens, out_features)
         # to figure out total tokens.
-        dy_2d = grad_output.reshape(-1, self.out_features)
-        total_tokens = dy_2d.size(0)
+        prequantized_mxfp8_grad = (
+            with_quantized_compute
+            and isinstance(grad_output, GroupedTensor)
+            and isinstance(ctx.grad_output_quantizers[0], MXFP8Quantizer)
+            and isinstance(grad_output.quantizer, MXFP8Quantizer)
+        )
+        if prequantized_mxfp8_grad:
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, out_features) layout; just validate the shape.
+            if grad_output.dim() != 2 or grad_output.size(-1) != self.out_features:
+                raise ValueError(
+                    "GroupedTensor grad output must have shape (total_tokens, "
+                    f"{self.out_features}), but got {tuple(grad_output.size())}."
+                )
+            dy_2d = None
+            total_tokens = grad_output.size(0)
+        else:
+            dy_2d = grad_output.reshape(-1, self.out_features)
+            total_tokens = dy_2d.size(0)
 
         # Build the grad_output GroupedTensor.
         # Optionally get dbias is fusion available with bgrad_group_quantize
@@ -1700,7 +1758,34 @@ def _fuser_backward_grouped_tensor(
             fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or (
                 isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad
             )
-            if has_bias and not self._scale_bias and fuse_bgrad:
+            if prequantized_mxfp8_grad:
+                # Rowwise-only MXFP8 grad output (e.g. FP8 token dispatch): reuse the
+                # rowwise data for the dgrad GEMM and manufacture the columnwise copy
+                # for wgrad. Bias grads are reduced from the dequantized grad below,
+                # which is only kept when there is a bias.
+                grouped_dy = grad_output.copy()
+                if ctx.weight_requires_grad:
+                    dy_2d = tex.group_requantize_columnwise_and_swizzle_rowwise_(
+                        grouped_dy,
+                        make_columnwise_gemm_quantizer(grad_output_quantizer),
+                        num_groups,
+                        split_sizes,
+                        TE_DType[dtype],
+                        tensor_offsets=base_split_offsets * self.out_features,
+                        return_dequantized=has_bias,
+                    )
+                else:
+                    # No wgrad, so no columnwise copy is needed. Dequantize before
+                    # swizzling: dequantization reads the unswizzled rowwise scales.
+                    dy_2d = (
+                        tex.group_dequantize(grouped_dy, TE_DType[dtype]).rowwise_data.view(
+                            total_tokens, self.out_features
+                        )
+                        if has_bias
+                        else None
+                    )
+                    tex.grouped_swizzle_for_gemm(grouped_dy, rowwise=True, columnwise=False)
+            elif has_bias and not self._scale_bias and fuse_bgrad:
                 grouped_dy, dbias_packed = tex.bgrad_group_quantize(
                     dy_2d, grad_output_quantizer, num_groups, split_sizes
                 )
@@ -1735,7 +1820,8 @@ def _fuser_backward_grouped_tensor(
                     offsets=base_split_offsets,
                 )
             elif dbias_packed is None:
-                # BF16/FP16 path
+                # BF16/FP16 and pre-quantized MXFP8 paths, neither of which fuses dbias
+                # into a quantize kernel.
                 dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups)
             if self.single_grouped_bias:
                 final_bias_grads = [dbias_packed.to(dtype=dtype)]