[PyTorch] Add optional caller-provided output/grad-input buffers to GroupedLinear and fused grouped MLP#3161
Conversation
Greptile SummaryThis PR adds optional caller-provided preallocated output and grad-input buffers to
Confidence Score: 5/5The change is safe to merge; the new buffer-passing paths are well-validated and the GEMM kernel integration is correct and fully additive. All three code paths (module API, ops legacy, ops GroupedTensor) validate the caller buffer before any GEMM work begins. The MXFP8 in-place kernel path uses a correctly structured as_strided view matching the kernel's native output layout. OperationFuser normalises basic_op_kwargs=None to empty-dict lists before reaching fused op forward methods. Tests cover alias, bit-exact match, padded-tail sentinel, and wrong-shape rejection across all paths. No files require special attention beyond the two style notes on duplicate validation logic and the error message in _resolve_op_kwargs. Important Files Changed
Sequence Diagram%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
participant Caller
participant Sequential
participant OperationFuser
participant GroupedLinear
participant GEMM Kernel
Caller->>Sequential: "forward(x, split_sizes, op_kwargs={op: {output: buf}})"
Sequential->>Sequential: _resolve_op_kwargs(op_kwargs)
Sequential->>OperationFuser: "__call__(*xs, basic_op_kwargs=[{output: buf}])"
Note over OperationFuser: normalises None to empty-dict list
OperationFuser->>GroupedLinear: fuser_forward(basic_op_kwargs)
GroupedLinear->>GroupedLinear: validate_or_alloc_output(buf, shape, dtype, device)
GroupedLinear->>GEMM Kernel: GEMM writes into buf in-place
GEMM Kernel-->>GroupedLinear: result in buf
GroupedLinear-->>Caller: y aliases buf, no copy
%%{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 GroupedLinear
participant GEMM Kernel
Caller->>Sequential: "forward(x, split_sizes, op_kwargs={op: {output: buf}})"
Sequential->>Sequential: _resolve_op_kwargs(op_kwargs)
Sequential->>OperationFuser: "__call__(*xs, basic_op_kwargs=[{output: buf}])"
Note over OperationFuser: normalises None to empty-dict list
OperationFuser->>GroupedLinear: fuser_forward(basic_op_kwargs)
GroupedLinear->>GroupedLinear: validate_or_alloc_output(buf, shape, dtype, device)
GroupedLinear->>GEMM Kernel: GEMM writes into buf in-place
GEMM Kernel-->>GroupedLinear: result in buf
GroupedLinear-->>Caller: y aliases buf, no copy
Reviews (9): Last reviewed commit: "Merge branch 'main' into pyt-gg-w-symm" | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| output: Optional[torch.Tensor] = None, | ||
| grad_input: Optional[torch.Tensor] = None, |
There was a problem hiding this comment.
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)There was a problem hiding this comment.
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.
| # when the value above was materialized into it; otherwise it is copied. | ||
| if output_buffer is not None: | ||
| output_buffer = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device) | ||
| if output_buffer.data_ptr() != fc2_out.data_ptr(): |
There was a problem hiding this comment.
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?
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
Hi, the cuDNN FE PR is merged. I will follow up to upgrade the cuDNN FE submodule in TE.
There was a problem hiding this comment.
Made a PR to upgrade cuDNN FE commit in TE submodule #3179
phu0ngng
changed the title
…ar module and fusible ops Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
for more information, see https://pre-commit.ci
…ping Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
…uffers, eliminating the D2D copy + cleanup Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
|
pipeline 57614140 |
Signed-off-by: YangFei1990 <feiw@nvidia.com>
Signed-off-by: YangFei1990 <feiw@nvidia.com>
|
/te-ci pytorch |
Signed-off-by: YangFei1990 <feiw@nvidia.com>
|
/te-ci pytorch |
1 similar comment
|
/te-ci pytorch |
3 participants
Description
Let callers have an option to supply preallocated buffers for the forward output and backward input gradient of
GroupedLinearand 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'scombine, eliminating the D2D copy from the token buffer into TE EP's internal staging buffer.Type of change
Changes
GroupedLinear.forward: optionalout/dgrad_outbuffers, wired through the legacy and GroupedTensor paths (incl. FP8). Validated via a sharedvalidate_or_alloc_outputhelper (shape/dtype/device/contiguous/no-grad).Sequential.forwardtakes anop_kwargsmapping (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.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, andValueErroron shape mismatch.test_grouped_linear_caller_buffers(ops) andtest_grouped_mlp_caller_buffers: verify the fused op runs, the output aliases the buffer, and grads match internal allocation.Checklist: