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

[phu0ngng]

Description

Let callers have an option to supply preallocated buffers for the forward output and backward input gradient of GroupedLinear and the fused grouped MLP, instead of the module/op allocating them internally. The caller can pass a symmetric-memory buffer that is fed directly into TE EP's combine, eliminating the D2D copy from the token buffer into TE EP's internal staging buffer.

Type of change

• Documentation change
• Bug fix
• New feature (non-breaking change which adds functionality)
• Breaking change
• Infra/Build change
• Code refactoring

Changes

• GroupedLinear.forward: optional out / dgrad_out buffers, wired through the legacy and GroupedTensor paths (incl. FP8). Validated via a shared validate_or_alloc_output helper (shape/dtype/device/contiguous/no-grad).
• Fusible ops: Sequential.forward takes an op_kwargs mapping (keyed by module or index) that routes per-op kwargs to the target op, used to pass output/grad-input buffers to grouped linear / grouped MLP.
• Fused grouped MLP: the GEMMs write the output and dgrad directly into the caller buffers for both nvfp4 and MXFP8, with no copy. The GEMM writes only the packed valid rows (per m_splits); any padded tail is left untouched. Buffers may be supplied independently.

Dependency

The in-place MXFP8 path requires a cuDNN frontend change that exposes an optional output tensor on the grouped-GEMM quant wrapper: NVIDIA/cudnn-frontend#338.

Testing

• test_grouped_linear_caller_output_buffers: covers out/dgrad_out/both across the legacy and GroupedTensor paths _ buffer aliasing, bit-exact match vs. internal allocation, untouched padded tail, and ValueError on shape mismatch.
• test_grouped_linear_caller_buffers (ops) and test_grouped_mlp_caller_buffers: verify the fused op runs, the output aliases the buffer, and grads match internal allocation.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR adds optional caller-provided output and grad-input buffers to GroupedLinear (both the nn.Module and BasicOperation variants) and the fused MXFP8/NVFP4 GroupedMLP, enabling zero-copy integration with symmetric-memory communication buffers in Expert Parallelism workflows.

• GroupedLinear.forward gains out/dgrad_out parameters, validated through a shared _validate_or_alloc_output helper; the ops-layer equivalent uses validate_or_alloc_output in _common.py.
• Sequential.forward gains op_kwargs (keyed by module or index) routed through a new _resolve_op_kwargs helper that maps per-op kwargs onto the correct slot in each OperationFuser's basic-op list.
• The fused MXFP8 GEMM path threads the caller buffer into the cuDNN kernel via as_strided so the kernel writes directly into the flat 2-D buffer, avoiding any D2D copy.

Confidence Score: 4/5

The buffer-aliasing feature is well-structured and validated at the boundary, but _OperationFuserAutogradFunction.forward mutates the caller's output buffer by calling requires_grad_(True) on it, which would break reuse of the same buffer across training iterations — the primary motivation for this feature.

The new validate_or_alloc_output helper, the _resolve_op_kwargs routing, and the as_strided kernel trick are each implemented correctly for a single forward+backward pass. The core concern is that _OperationFuserAutogradFunction.forward calls x.requires_grad_(True) where x is the caller-provided buffer. Because this is an in-place Python mutation, output_buffer.requires_grad remains True after the call returns, and the next iteration's validate_or_alloc_output would raise ValueError. Tests cover one-shot usage only and do not exercise iteration reuse, so the defect is not currently caught.

transformer_engine/pytorch/ops/fused/grouped_mlp.py (output buffer aliasing with the fuser's requires_grad mutation) and tests/pytorch/test_grouped_mlp.py / test_fusible_ops.py (missing multi-iteration buffer reuse coverage).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds out/dgrad_out parameters to GroupedLinear.forward and threads them through the grouped-tensor and legacy backward paths via _validate_or_alloc_output.
transformer_engine/pytorch/ops/_common.py	Adds OUTPUT_BUFFER_KEY/GRAD_INPUT_BUFFER_KEY string constants and a reusable validate_or_alloc_output helper used throughout the ops stack.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Wires out_buffer/dgrad_out from basic_op_kwargs into both legacy and grouped-tensor forward/backward paths; dgrad_out is retrieved from ctx with getattr(..., None) for safety.
transformer_engine/pytorch/ops/fused/grouped_mlp.py	Pulls output/grad-input buffers from basic_op_kwargs and writes GEMM results directly into them via as_strided; the caller buffer's requires_grad flag is mutated by _OperationFuserAutogradFunction.forward, breaking buffer reuse across iterations.
transformer_engine/pytorch/ops/sequential.py	Adds op_kwargs parameter to Sequential.forward and _resolve_op_kwargs helper that maps module/index keys to per-basic-op kwargs lists; integer key out-of-range yields an undifferentiated IndexError.
tests/pytorch/test_fusible_ops.py	Adds test_grouped_linear_caller_buffers covering forward aliasing, shape rejection, and a single-pass backward; does not test buffer reuse across iterations.
tests/pytorch/test_grouped_linear.py	Adds test_grouped_linear_caller_output_buffers with parametrized out/dgrad_out/both coverage, padded-tail sentinel checks, and a shape-mismatch rejection test.
tests/pytorch/test_grouped_mlp.py	Adds test_grouped_mlp_caller_buffers for the MXFP8 fused op; verifies aliasing and grad match for a single forward+backward pass, without testing iteration reuse.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller
    participant Sequential
    participant OperationFuser
    participant GroupedMLP_CuTeGEMMBase
    participant cuDNN_Kernel

    Caller->>Sequential: forward(x, split_sizes, op_kwargs)
    Sequential->>Sequential: _resolve_op_kwargs(op_kwargs)
    Sequential->>OperationFuser: __call__(x, basic_op_kwargs)
    OperationFuser->>GroupedMLP_CuTeGEMMBase: fuser_forward(basic_op_ctxs, x, basic_op_kwargs)
    GroupedMLP_CuTeGEMMBase->>GroupedMLP_CuTeGEMMBase: validate_or_alloc_output(output_buffer)
    GroupedMLP_CuTeGEMMBase->>cuDNN_Kernel: "grouped_gemm_quant_kernel(d_tensor=out_buf.as_strided(...))"
    cuDNN_Kernel-->>GroupedMLP_CuTeGEMMBase: writes directly into out_buf storage
    GroupedMLP_CuTeGEMMBase-->>OperationFuser: "fc2_out = output_buffer"
    OperationFuser-->>Caller: y (aliases out_buf)

    Caller->>OperationFuser: y.backward(dy)
    OperationFuser->>GroupedMLP_CuTeGEMMBase: fuser_backward(fc1_ctx)
    GroupedMLP_CuTeGEMMBase->>cuDNN_Kernel: "fc1_dgrad_kernel(d_tensor=dgrad_buf.as_strided(...))"
    cuDNN_Kernel-->>GroupedMLP_CuTeGEMMBase: writes directly into dgrad_buf storage
    GroupedMLP_CuTeGEMMBase-->>Caller: "grad_input = dgrad_buf"

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller
    participant Sequential
    participant OperationFuser
    participant GroupedMLP_CuTeGEMMBase
    participant cuDNN_Kernel

    Caller->>Sequential: forward(x, split_sizes, op_kwargs)
    Sequential->>Sequential: _resolve_op_kwargs(op_kwargs)
    Sequential->>OperationFuser: __call__(x, basic_op_kwargs)
    OperationFuser->>GroupedMLP_CuTeGEMMBase: fuser_forward(basic_op_ctxs, x, basic_op_kwargs)
    GroupedMLP_CuTeGEMMBase->>GroupedMLP_CuTeGEMMBase: validate_or_alloc_output(output_buffer)
    GroupedMLP_CuTeGEMMBase->>cuDNN_Kernel: "grouped_gemm_quant_kernel(d_tensor=out_buf.as_strided(...))"
    cuDNN_Kernel-->>GroupedMLP_CuTeGEMMBase: writes directly into out_buf storage
    GroupedMLP_CuTeGEMMBase-->>OperationFuser: "fc2_out = output_buffer"
    OperationFuser-->>Caller: y (aliases out_buf)

    Caller->>OperationFuser: y.backward(dy)
    OperationFuser->>GroupedMLP_CuTeGEMMBase: fuser_backward(fc1_ctx)
    GroupedMLP_CuTeGEMMBase->>cuDNN_Kernel: "fc1_dgrad_kernel(d_tensor=dgrad_buf.as_strided(...))"
    cuDNN_Kernel-->>GroupedMLP_CuTeGEMMBase: writes directly into dgrad_buf storage
    GroupedMLP_CuTeGEMMBase-->>Caller: "grad_input = dgrad_buf"

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into pyt-gg-w-symm" | Re-trigger Greptile}

[timmoon10]

I'm skeptical of this API because it looks general, but is actually only supported with grouped linear and grouped MLP. How about something like the following:

def forward(
    self,
    input: torch.Tensor,
    *extra_inputs: torch.Tensor,
    extra_kwargs: Optional[Sequence[Optional[dict[str, Any]]]] = None,
):
    ...

    if extra_kwargs is None:
        extra_kwargs = [None] * len(self)

    # Figure out what extra_kwargs go to each module group, and fail if a non-None extra_kwargs goes to a non-basic op

    for module_group in self._module_groups:
            xs = module_group(*xs, basic_op_kwargs=module_group_basic_op_kwargs)

[phu0ngng]

Thanks, Tim. Good call. The output and grad_input kwargs overpromised generality. Building on your suggestion, I went with a generic per op kwargs pass through, but keyed by module or index instead of a positional list:

seq(x, split_sizes, probs, split_sizes,
    op_kwargs={fc1: {GRAD_INPUT_BUFFER_KEY: g}, fc2: {OUTPUT_BUFFER_KEY: o}})

Sequential resolves each key to its group and basic op slot and forwards it as basic_op_kwargs, with no new dispatch path, and raises if a key targets a non-fusible op. This keeps the container kwarg-agnostic, so the buffers are now just op specific kwargs, while avoiding None padding and positional coupling.

Let me know what you think about this new design.

[YangFei1990]

Under what case this condition can be false? It seems to be it is always true, I do not see we feed the user provided buffer to the kernel?

[phu0ngng]

You were right that it was always true for MXFP8, which uses CuTe GEMM.

I pushed the latest commit in which I reworked it so that it does not trigger extra copies: the caller buffer is passed straight into the CuTe GEMM via d_tensor, so the kernel writes output/dgrad in place for both nvfp4 and MXFP8. This needs a cuDNN FE change I made a PR for: NVIDIA/cudnn-frontend#338 (until it lands and the bundled FE is bumped, the in-place MXFP8 path can't run). So for testing, we will need to check out this cuDNN FE manually.

[phu0ngng]

Hi, the cuDNN FE PR is merged. I will follow up to upgrade the cuDNN FE submodule in TE.

[phu0ngng]

Made a PR to upgrade cuDNN FE commit in TE submodule #3179

[phu0ngng]

pipeline 57614140

[greptile-apps]

Output buffer requires_grad mutated by fuser forward

_OperationFuserAutogradFunction.forward calls x.requires_grad_(True) on the final output tensor (x) when gradient computation is enabled. Because x is the caller's output_buffer (set via fc2_out = output_buffer in the MXFP8 branch), this permanently sets output_buffer.requires_grad = True. On any subsequent forward call that provides the same buffer, validate_or_alloc_output raises ValueError: "Output buffer must not require gradient." — which breaks the intended use case of reusing a symmetric-memory buffer across training iterations. The same applies to fc1_ctx.dgrad_out on the backward side. The fuser's line x.requires_grad_(fuser.first_op_requiring_backward < fuser._num_basic_ops) in fuser.py is the mutation site.