Improve device-init grouped linear module with single grouped weight support

tests/pytorch/test_grouped_mlp.py

@@ -23,6 +23,7 @@
     OUTPUT_BUFFER_KEY,
     GRAD_INPUT_BUFFER_KEY,
 )
+from transformer_engine.pytorch.utils import get_device_compute_capability
 from transformer_engine.pytorch import (
     QuantizedTensor,
     Float8CurrentScalingQuantizer,
@@ -129,6 +130,28 @@ def maybe_skip_quantization(
             pytest.skip("NVFP4 quantization is only supported with BF16 data")
 
 
+def grouped_tensor_path_supported(
+    *,
+    quantized_compute: bool,
+    quantization: Optional[str],
+    dtype: torch.dtype,
+    single_grouped_weight: bool,
+) -> bool:
+    """Mirror GroupedLinear's native grouped-tensor backend selection."""
+    compute_capability = get_device_compute_capability()
+    if not (9, 0) <= compute_capability <= (11, 0):
+        return False
+    if not quantized_compute:
+        return dtype in (torch.bfloat16, torch.float16)
+    if quantization == "fp8_current_scaling":
+        return compute_capability >= (10, 0) or tex.get_cublasLt_version() >= 130500
+    if quantization == "mxfp8":
+        return compute_capability >= (10, 0)
+    if quantization == "nvfp4_rht":
+        return compute_capability >= (10, 0) and not single_grouped_weight
+    return False
+
+
 @torch.no_grad()
 def make_reference_and_test_tensors(
     shape: int | Iterable[int],
@@ -232,6 +255,26 @@ def make_reference_and_test_tensors(
 class TestGroupedLinearOp:
     """Tests for advanced features with grouped linear basic op"""
 
+    def test_single_grouped_bias_uses_registered_packed_storage(self, monkeypatch) -> None:
+        """The grouped bias compute view must alias the registered trainable parent."""
+        monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+        op = te.ops.GroupedLinear(
+            2,
+            128,
+            128,
+            bias=True,
+            device="cuda",
+            dtype=torch.bfloat16,
+            single_grouped_bias=True,
+        )
+
+        bias_packed = op._get_packed_bias_tensor(torch.bfloat16)
+
+        assert bias_packed.shape == (op.num_groups, op.out_features)
+        assert bias_packed.untyped_storage().data_ptr() == op.bias.rowwise_data.data_ptr()
+        assert op.bias.requires_grad
+        assert dict(op.named_parameters())["bias"] is op.bias
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("dtype", _dtypes)
     @pytest.mark.parametrize("quantization", _quantization_list)
@@ -292,6 +335,14 @@ def test_grouped_linear(
 
         if single_grouped_bias and not bias:
             pytest.skip("single_grouped_bias requires bias=True")
+        if (single_grouped_weight or single_grouped_bias) and not grouped_tensor_path_supported(
+            quantized_compute=quantized_compute,
+            quantization=quantization,
+            dtype=dtype,
+            single_grouped_weight=single_grouped_weight,
+        ):
+            # Single grouped parameters intentionally have no split-quantize fallback.
+            pytest.skip("Single grouped parameters require the native grouped-tensor path")
         if single_grouped_weight and quantized_weight and quantization in ("fp8_delayed_scaling"):
             pytest.skip(
                 "single_grouped_weight does not support FP8 delayed scaling "

tests/pytorch/test_grouped_tensor.py

@@ -1204,6 +1204,147 @@ def test_group_dequantize_cudagraph_capturable(self) -> None:
         for exp, got in zip(expected_tensors, static_tensors):
             assert torch.equal(got, exp)
 
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_group_quantize_reuses_destination_and_honors_noop(self) -> None:
+        """Verify the in-place and graph-safe contracts of grouped MXFP8 quantization.
+
+        A captured training graph records the addresses of all MXFP8 weight storage, not just
+        the Python ``GroupedTensor`` object. Re-quantizing an updated BF16 weight must therefore
+        refresh the existing FP8 payloads and scales without replacing any allocation. During
+        graph replay, ``noop_flag`` must additionally allow the captured quantization kernel to
+        execute as a no-op on microbatches that should reuse the cached weight.
+        """
+        num_tensors = 2
+        # Treat two contiguous [256, 256] matrices as one grouped BF16 source. These dimensions
+        # satisfy the MXFP8 block-alignment requirements in both rowwise and columnwise layouts.
+        shape = (num_tensors * 256, 256)
+        quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)
+
+        # Forward GEMM consumes rowwise weight storage, while backward dgrad consumes columnwise
+        # storage. Both representations and both scale tensors must be updated and kept stable.
+        quantizer.set_usage(rowwise=True, columnwise=True)
+        # Store the quantized output in GEMM-ready (swizzled) scale layout. This makes the test
+        # cover the exact destination format used by cached MXFP8 model weights.
+        quantizer.optimize_for_gemm = True
+
+        # The first call owns allocation: it creates the persistent destination that subsequent
+        # optimizer updates and CUDA graph replays must modify in place.
+        source = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+        output = tex.group_quantize(source, quantizer, num_tensors, None)
+
+        # CUDA graphs retain these raw addresses. Object identity alone is insufficient because
+        # replacing one internal payload or scale allocation would leave the graph with a stale
+        # pointer even if ``output`` remained the same Python object.
+        pointers = (
+            output.rowwise_data.data_ptr(),
+            output.columnwise_data.data_ptr(),
+            output.scale_inv.data_ptr(),
+            output.columnwise_scale_inv.data_ptr(),
+        )
+
+        # Eager in-place update: quantize new BF16 values into the previously allocated output.
+        # An independently allocated quantization provides the numerical reference.
+        updated_source = source + 1
+        updated = tex.group_quantize(
+            updated_source,
+            quantizer,
+            num_tensors,
+            None,
+            output=output,
+        )
+        reference = tex.group_quantize(updated_source, quantizer, num_tensors, None)
+
+        # The API must return the caller-provided destination and preserve every captured address.
+        assert updated is output
+        assert pointers == (
+            output.rowwise_data.data_ptr(),
+            output.columnwise_data.data_ptr(),
+            output.scale_inv.data_ptr(),
+            output.columnwise_scale_inv.data_ptr(),
+        )
+
+        # In-place quantization must still produce exactly the same FP8 bytes and scales as a
+        # normal allocating call, for both GEMM orientations.
+        assert torch.equal(output.rowwise_data, reference.rowwise_data)
+        assert torch.equal(output.columnwise_data, reference.columnwise_data)
+        assert torch.equal(output.scale_inv, reference.scale_inv)
+        assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv)
+
+        # Eager no-op: a nonzero device flag tells the kernel to preserve the cached destination.
+        # Clone all four buffers so this checks payloads and metadata independently.
+        old_buffers = (
+            output.rowwise_data.clone(),
+            output.columnwise_data.clone(),
+            output.scale_inv.clone(),
+            output.columnwise_scale_inv.clone(),
+        )
+        noop = torch.ones(1, dtype=torch.float32, device="cuda")
+
+        # The source is deliberately different. If the kernel ignores ``noop_flag``, at least
+        # one payload or scale tensor below will change and expose the failure.
+        tex.group_quantize(
+            updated_source * 2,
+            quantizer,
+            num_tensors,
+            None,
+            noop_flag=noop,
+            output=output,
+        )
+        assert torch.equal(output.rowwise_data, old_buffers[0])
+        assert torch.equal(output.columnwise_data, old_buffers[1])
+        assert torch.equal(output.scale_inv, old_buffers[2])
+        assert torch.equal(output.columnwise_scale_inv, old_buffers[3])
+
+        # Capture one fixed launch. ``graph_source``, ``noop``, and ``output`` keep stable device
+        # addresses; replay changes only their contents. This is how weight caching works when
+        # different microbatches replay the same CUDA graph.
+        graph_source = updated_source.clone()
+        # zero means "perform quantization"; nonzero means "leave destination untouched".
+        noop.zero_()
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            tex.group_quantize(
+                graph_source,
+                quantizer,
+                num_tensors,
+                None,
+                noop_flag=noop,
+                output=output,
+            )
+
+        # Replay with new source contents and noop=0. The captured kernel must refresh the same
+        # destination allocations, and the result must equal a fresh eager quantization.
+        graph_source.copy_(source - 1)
+        noop.zero_()
+        graph.replay()
+        torch.cuda.synchronize()
+        reference = tex.group_quantize(graph_source, quantizer, num_tensors, None)
+        assert torch.equal(output.rowwise_data, reference.rowwise_data)
+        assert torch.equal(output.columnwise_data, reference.columnwise_data)
+        assert torch.equal(output.scale_inv, reference.scale_inv)
+        assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv)
+
+        # Snapshot the successfully refreshed cache before testing the captured no-op branch.
+        old_buffers = (
+            output.rowwise_data.clone(),
+            output.columnwise_data.clone(),
+            output.scale_inv.clone(),
+            output.columnwise_scale_inv.clone(),
+        )
+
+        # Replay the exact same captured graph with different source data but noop=1. The launch
+        # still occurs, which is required for CUDA graph structural consistency, but the kernel
+        # must not mutate any cached FP8 payload or scale buffer.
+        graph_source.copy_(source + 2)
+        noop.fill_(1)
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(output.rowwise_data, old_buffers[0])
+        assert torch.equal(output.columnwise_data, old_buffers[1])
+        assert torch.equal(output.scale_inv, old_buffers[2])
+        assert torch.equal(output.columnwise_scale_inv, old_buffers[3])
+
     @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"])
     @pytest.mark.parametrize("direction", ["rowwise", "columnwise"])
     @pytest.mark.skipif(

transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh

@@ -572,9 +572,18 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
     const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise);
 
     const size_t tensor_base = current_block.tensor_base;
-    const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1)
-                                              ? static_cast<size_t>(offsets_ptr[tensor_id])
-                                              : tensor_base;
+    size_t tensor_base_for_scales = tensor_base;
+    size_t tensor_rows_for_scales = rows;
+    if constexpr (WITH_GEMM_SWIZZLED_SCALES && SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) {
+      // The payload is one tall tensor, but GEMM scales are swizzled independently per member.
+      // Uniform groups omit first_dims/tensor_offsets, so derive the member geometry directly.
+      tensor_rows_for_scales = first_logical_dim / num_tensors;
+      tensor_base_for_scales = tensor_id * tensor_rows_for_scales * cols;
+    }
+    if constexpr (WITH_GEMM_SWIZZLED_SCALES &&
+                  SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) {
+      tensor_base_for_scales = static_cast<size_t>(offsets_ptr[tensor_id]);
+    }
     const size_t block_id_Y = current_block.block_id_Y;
     const size_t block_id_X = current_block.block_id_X;
     const size_t block_offset_Y = current_block.block_offset_Y;
@@ -680,8 +689,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
         process_colwise_stage<IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType, OType, ROWWISE_SCALING,
                               WITH_GEMM_SWIZZLED_SCALES>(
             buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise,
-            scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr,
-            sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise);
+            scale_stride_colwise, tensor_base_for_scales, tensor_rows_for_scales, cols, sIn_ptr,
+            sActIn_ptr, sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise);
       }
 
       if constexpr (ROWWISE_SCALING) {

transformer_engine/pytorch/csrc/extensions.h

@@ -350,7 +350,7 @@ py::object dequantize(const py::handle &input, DType otype);
 py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,
                           std::optional<at::Tensor> first_dims, std::optional<at::Tensor> last_dims,
                           std::optional<at::Tensor> tensor_offsets,
-                          std::optional<at::Tensor> noop_flag);
+                          std::optional<at::Tensor> noop_flag, const py::object &output);
 
 py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
                                           const size_t num_tensors,

transformer_engine/pytorch/csrc/extensions/cast.cpp

@@ -276,7 +276,7 @@ void compute_grouped_fp8_current_scaling_amax_and_scale(
 py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,
                           std::optional<at::Tensor> first_dims, std::optional<at::Tensor> last_dims,
                           std::optional<at::Tensor> tensor_offsets,
-                          std::optional<at::Tensor> noop_flag) {
+                          std::optional<at::Tensor> noop_flag, const py::object &output) {
   using namespace transformer_engine::pytorch::detail;
   init_extension();
 
@@ -305,11 +305,34 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const
                                         GetTransformerEngineDType(tensor.scalar_type()),
                                         std::vector<size_t>{static_cast<size_t>(tensor.numel())});
 
-  // Create output GroupedTensor.
-  auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor(
-      num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
-      py::reinterpret_borrow<py::object>(quantizer), first_dims, last_dims, tensor_offsets,
-      logical_first_dim, logical_last_dim);
+  // Create a GroupedTensor or reuse an existing destination. Reusing the destination is
+  // required for weight caching and CUDA graph replay because captured GEMMs retain the
+  // original data and scale pointers.
+  py::object grouped_output_py;
+  auto grouped_output_tensor_cpp = [&]() -> GroupedTensorWrapper {
+    if (!output.is_none()) {
+      NVTE_CHECK(!first_dims.has_value() && !last_dims.has_value() && !tensor_offsets.has_value(),
+                 "group_quantize: output reuse currently requires uniform tensor shapes.");
+      NVTE_CHECK(output.attr("num_tensors").cast<size_t>() == num_tensors,
+                 "group_quantize: output has a different number of tensors.");
+      NVTE_CHECK(output.attr("logical_shape").cast<std::vector<size_t>>() == logical_shape,
+                 "group_quantize: output has an incompatible logical shape.");
+      py::object output_quantizer = output.attr("quantizer");
+      NVTE_CHECK(!output_quantizer.is_none(),
+                 "group_quantize: output must have quantized storage.");
+      NVTE_CHECK(output_quantizer.is(quantizer),
+                 "group_quantize: output must have been created by the same quantizer.");
+      grouped_output_py = output;
+      return GroupedTensorFromPyTorchGroupedTensor(output);
+    }
+
+    auto result = quantizer_cpp->create_grouped_tensor(
+        num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
+        py::reinterpret_borrow<py::object>(quantizer), first_dims, last_dims, tensor_offsets,
+        logical_first_dim, logical_last_dim);
+    grouped_output_py = std::move(result.second);
+    return std::move(result.first);
+  }();
 
   // dispatch to scaling methods
   enum class GroupedQuantizationMode {

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -208,7 +208,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("group_quantize", transformer_engine::pytorch::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(),
-        py::arg("noop_flag") = py::none());
+        py::arg("noop_flag") = py::none(), py::arg("output") = py::none());
   transformer_engine::pytorch::bind_quantize_with_amax_extensions(m);
   m.def("group_dequantize", transformer_engine::pytorch::group_dequantize,
         "Dequantize group tensor", py::arg("input"), py::arg("otype"));