[PyTorch] Add optional caller-provided output/grad-input buffers to GroupedLinear and fused grouped MLP

tests/pytorch/test_fusible_ops.py

@@ -19,6 +19,10 @@
 import transformer_engine.common.recipe
 import transformer_engine.pytorch as te
 import transformer_engine.pytorch.ops as te_ops
+from transformer_engine.pytorch.ops._common import (
+    OUTPUT_BUFFER_KEY,
+    GRAD_INPUT_BUFFER_KEY,
+)
 from transformer_engine.pytorch._extra_state import UNSAFE_PICKLE_EXTRA_STATE_ENV
 
 from transformer_engine.pytorch.ops.fused import (
@@ -2196,6 +2200,86 @@ def test_grouped_linear(
             if bias:
                 assert_close_grads(getattr(op, f"bias{group_idx}"), bs_ref[group_idx], **tols)
 
+    def test_grouped_linear_caller_buffers(
+        self,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+    ) -> None:
+        """Caller output/grad_input buffers routed to the last/first op of a Sequential.
+
+        The chain has two GroupedLinears with distinct inner dims, so ``output``
+        can only fit the last op's output and ``grad_input`` only the first op's
+        dgrad -- a mis-route would fail the buffer's shape check.
+        """
+        group_size = 3
+        in_features, hidden, out_features = 128, 256, 64
+        split_sizes = torch.tensor([128, 256, 128], dtype=torch.int32, device=device)
+        num_tokens = int(split_sizes.sum())
+
+        torch.manual_seed(1234)
+        x = (0.1 * torch.randn(num_tokens, in_features, device=device)).to(dtype)
+        dy = (0.1 * torch.randn(num_tokens, out_features, device=device)).to(dtype)
+
+        def build() -> te_ops.Sequential:
+            fc1 = te_ops.GroupedLinear(
+                group_size, in_features, hidden, bias=False, device=device, dtype=dtype
+            )
+            fc2 = te_ops.GroupedLinear(
+                group_size, hidden, out_features, bias=False, device=device, dtype=dtype
+            )
+            return te_ops.Sequential(fc1, fc2)
+
+        # Reference: internal allocation.
+        model_ref = build()
+        x_ref = x.detach().clone().requires_grad_(True)
+        y_ref = model_ref(x_ref, split_sizes, split_sizes)
+        y_ref.backward(dy)
+
+        # Caller-provided buffers, same weights as the reference.
+        model = build()
+        with torch.no_grad():
+            for op_idx in range(2):
+                for i in range(group_size):
+                    getattr(model[op_idx], f"weight{i}").copy_(
+                        getattr(model_ref[op_idx], f"weight{i}")
+                    )
+        sentinel = 7.0
+        out_buf = torch.full((num_tokens, out_features), sentinel, dtype=dtype, device=device)
+        dgrad_buf = torch.full((num_tokens, in_features), sentinel, dtype=dtype, device=device)
+        x_test = x.detach().clone().requires_grad_(True)
+        y = model(
+            x_test,
+            split_sizes,
+            split_sizes,
+            op_kwargs={
+                model[0]: {GRAD_INPUT_BUFFER_KEY: dgrad_buf},
+                model[1]: {OUTPUT_BUFFER_KEY: out_buf},
+            },
+        )
+
+        # Forward output aliases the last op's output buffer with no copy.
+        assert y.data_ptr() == out_buf.data_ptr()
+        torch.testing.assert_close(y, y_ref, rtol=0, atol=0)
+
+        y.backward(dy)
+
+        # grad_input written into the first op's dgrad buffer.
+        assert not torch.all(dgrad_buf == sentinel)
+        torch.testing.assert_close(dgrad_buf, x_ref.grad, rtol=0, atol=0)
+        torch.testing.assert_close(x_test.grad, x_ref.grad, rtol=0, atol=0)
+
+        # A buffer whose shape does not match the output is rejected.
+        bad = torch.empty(num_tokens + 1, out_features, dtype=dtype, device=device)
+        with pytest.raises(ValueError):
+            model_bad = build()
+            model_bad(
+                x.detach(),
+                split_sizes,
+                split_sizes,
+                op_kwargs={model_bad[1]: {OUTPUT_BUFFER_KEY: bad}},
+            )
+
     @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128)))
     @pytest.mark.parametrize("input_requires_grad", (False, True))
     @pytest.mark.parametrize("scales_requires_grad", (False, True))

tests/pytorch/test_grouped_linear.py

@@ -1710,6 +1710,102 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk
     grouped_linear.backward_dw()
 
 
+@pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"])
+@pytest.mark.parametrize("supply", ["out", "dgrad_out", "both"])
+def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatch):
+    """Caller-provided forward out and/or backward dgrad_out buffers.
+
+    Checks that a supplied buffer is written in place (bit-for-bit vs internal allocation)
+    and, on the fused path with a padded input, that only the valid rows are touched.
+    """
+    if use_fused_path:
+        device_capability = torch.cuda.get_device_capability()
+        if not (9, 0) <= device_capability <= (11, 0):
+            pytest.skip(
+                "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell"
+                " (SM10x and SM110)."
+            )
+        cublaslt_version = tex.get_cublasLt_version()
+        if device_capability < (10, 0) and cublaslt_version < 130400:
+            pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
+        if cublaslt_version < 130300:
+            pytest.skip("Grouped GEMM requires cuBLAS 13.3+.")
+
+    monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if use_fused_path else "0")
+    give_out = supply in ("out", "both")
+    give_dgrad = supply in ("dgrad_out", "both")
+
+    dtype = torch.bfloat16
+    num_gemms = 3
+    in_features = 128
+    out_features = 128
+    m_splits_list = [64, 96, 80]
+    valid_tokens = sum(m_splits_list)  # 240
+    # The fused path supports a padded input; the legacy path requires tight packing.
+    num_rows = valid_tokens + (80 if use_fused_path else 0)
+    m_splits = (
+        torch.tensor(m_splits_list, dtype=torch.int64, device="cuda")
+        if use_fused_path
+        else m_splits_list
+    )
+
+    torch.manual_seed(1234)
+    x_base = (0.1 * torch.randn(num_rows, in_features, device="cuda")).to(dtype)
+    dy = torch.zeros(num_rows, out_features, dtype=dtype, device="cuda")
+    dy[:valid_tokens] = (0.1 * torch.randn(valid_tokens, out_features, device="cuda")).to(dtype)
+
+    grouped_linear = GroupedLinear(
+        num_gemms,
+        in_features,
+        out_features,
+        bias=False,
+        params_dtype=dtype,
+        device="cuda",
+    )
+
+    # Reference: internal allocation.
+    x_ref = x_base.detach().clone().requires_grad_(True)
+    y_ref = grouped_linear(x_ref, m_splits)
+    y_ref.backward(dy)
+
+    # Caller-provided buffers with a sentinel-filled tail (only the requested ones).
+    sentinel = 7.0
+    out_buf = (
+        torch.full((num_rows, out_features), sentinel, dtype=dtype, device="cuda")
+        if give_out
+        else None
+    )
+    dgrad_buf = (
+        torch.full((num_rows, in_features), sentinel, dtype=dtype, device="cuda")
+        if give_dgrad
+        else None
+    )
+    x = x_base.detach().clone().requires_grad_(True)
+    y = grouped_linear(x, m_splits, out=out_buf, dgrad_out=dgrad_buf)
+
+    if give_out:
+        # Forward output is the caller buffer itself (no copy); padded tail untouched.
+        assert y.data_ptr() == out_buf.data_ptr()
+        assert tuple(y.shape) == (num_rows, out_features)
+        assert torch.all(out_buf[valid_tokens:] == sentinel)
+    torch.testing.assert_close(y[:valid_tokens], y_ref[:valid_tokens], rtol=0, atol=0)
+
+    y.backward(dy)
+
+    if give_dgrad:
+        # dgrad written into the caller buffer; padded tail untouched.
+        assert torch.all(dgrad_buf[valid_tokens:] == sentinel)
+        torch.testing.assert_close(
+            dgrad_buf[:valid_tokens], x_ref.grad[:valid_tokens], rtol=0, atol=0
+        )
+    torch.testing.assert_close(x.grad[:valid_tokens], x_ref.grad[:valid_tokens], rtol=0, atol=0)
+
+    # A buffer whose row count does not match the input rows is rejected.
+    bad_out = torch.empty(num_rows + 1, out_features, dtype=dtype, device="cuda")
+    with pytest.raises(ValueError):
+        grouped_linear(x, m_splits, out=bad_out)
+
+
 @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4)
 def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
     """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct.

tests/pytorch/test_grouped_mlp.py

@@ -19,6 +19,10 @@
     _cudnn_frontend_supports_grouped_gemm_srelu,
     _cudnn_frontend_version_supported,
 )
+from transformer_engine.pytorch.ops._common import (
+    OUTPUT_BUFFER_KEY,
+    GRAD_INPUT_BUFFER_KEY,
+)
 from transformer_engine.pytorch import (
     QuantizedTensor,
     Float8CurrentScalingQuantizer,
@@ -1317,6 +1321,98 @@ def _run_case(single_grouped_weight: bool) -> tuple[torch.Tensor, ...]:
             torch.testing.assert_close(fc1_db_false, fc1_db_true, **bias_tols)
             torch.testing.assert_close(fc2_db_false, fc2_db_true, **bias_tols)
 
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_grouped_mlp_caller_buffers(
+        self,
+        *,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        group_size: int = 4,
+        hidden_size: int = 256,
+        split_alignment: int = 256,
+        glu_interleave_size: int = 32,
+    ) -> None:
+        """Caller-provided output/grad_input buffers on the fused MXFP8 grouped MLP."""
+        if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+
+        split_sizes = torch.tensor(
+            [split_alignment * (i + 1) for i in range(group_size)],
+            dtype=torch.int64,
+            device=device,
+        )
+        num_tokens = int(split_sizes.sum())
+        recipe = make_recipe("mxfp8")
+
+        x_base = torch.empty((num_tokens, hidden_size), device=device, dtype=dtype).uniform_(
+            -0.25, 0.25
+        )
+        probs_base = torch.empty((num_tokens,), device=device, dtype=dtype).uniform_(-0.25, 0.25)
+        dy_base = torch.empty((num_tokens, hidden_size), device=device, dtype=dtype).uniform_(
+            -0.25, 0.25
+        )
+        fc1_ws_base = [
+            torch.empty((2 * hidden_size, hidden_size), device=device, dtype=dtype).uniform_(
+                -0.25, 0.25
+            )
+            for _ in range(group_size)
+        ]
+        fc2_ws_base = [
+            torch.empty((hidden_size, hidden_size), device=device, dtype=dtype).uniform_(
+                -0.25, 0.25
+            )
+            for _ in range(group_size)
+        ]
+
+        def build() -> te.ops.Sequential:
+            with te.quantized_model_init(enabled=True, recipe=recipe):
+                fc1 = te.ops.GroupedLinear(
+                    group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype
+                )
+                scaled_act = te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size)
+                fc2 = te.ops.GroupedLinear(
+                    group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype
+                )
+                module = te.ops.Sequential(fc1, scaled_act, fc2)
+            with torch.no_grad():
+                for i in range(group_size):
+                    getattr(fc1, f"weight{i}").copy_(fc1_ws_base[i])
+                    getattr(fc2, f"weight{i}").copy_(fc2_ws_base[i])
+            return module
+
+        def run(module, *, output=None, grad_input=None):
+            x = x_base.detach().clone().requires_grad_(True)
+            probs = probs_base.detach().clone().requires_grad_(True)
+            op_kwargs = {}
+            if grad_input is not None:
+                op_kwargs[module[0]] = {GRAD_INPUT_BUFFER_KEY: grad_input}
+            if output is not None:
+                op_kwargs[module[2]] = {OUTPUT_BUFFER_KEY: output}
+            with te.autocast(enabled=True, recipe=recipe):
+                y = module(x, split_sizes, probs, split_sizes, op_kwargs=op_kwargs or None)
+            y.backward(dy_base)
+            return y, x.grad
+
+        # Reference: internal allocation.
+        y_ref, dx_ref = run(build())
+
+        # Caller-provided buffers.
+        module = build()
+        sentinel = 7.0
+        out_buf = torch.full((num_tokens, hidden_size), sentinel, dtype=dtype, device=device)
+        dgrad_buf = torch.full((num_tokens, hidden_size), sentinel, dtype=dtype, device=device)
+        y, _ = run(module, output=out_buf, grad_input=dgrad_buf)
+
+        # The fused op ran, the returned output aliases the buffer, and both buffers were written.
+        assert isinstance(
+            module._module_groups[0]._forward_ops[0][0],
+            te.ops.fused.GroupedMLP_CuTeGEMMGLU,
+        )
+        assert y.data_ptr() == out_buf.data_ptr()
+        assert not torch.all(dgrad_buf == sentinel)
+        torch.testing.assert_close(y, y_ref, rtol=0, atol=0)
+        torch.testing.assert_close(dgrad_buf, dx_ref, rtol=0, atol=0)
+
     @pytest.mark.parametrize("single_grouped_weight", (False, True))
     @pytest.mark.parametrize("delay_wgrad_compute", (False, True))
     @pytest.mark.parametrize("zero_out_wgrad", (False, True))