[PyTorch] Expose function to bulk-allocate tensors backed by the same buffer

[timmoon10]

Description

Allocating PyTorch tensors carries non-trivial CPU overhead, which becomes especially painful when allocating per-expert tensors in the grouped linear layer. This PR generalizes the bulk allocation approach used in the split-quantize functions (see #1793), which involves allocating a single large buffer and creating tensor subviews with at::from_blob. We expose a dedicated bulk-allocation function that is exposed to Python, refactor the split-quantize functions, and bulk-allocate wgrads in the grouped linear implementations.

This is incremental progress toward #2897. When I run the grouped MLP benchmark discussed in that issue, I see a runtime reduction of 150 us (for reference the backward pass takes ~2.1 ms).

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add dedicated function for bulk-allocating PyTorch tensors
• Bulk-allocate wgrad tensors in grouped linear implementations

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

[timmoon10]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR introduces a dedicated bulk_allocate C++ function that allocates a single contiguous CUDA buffer and returns per-tensor views via at::from_blob, eliminating the CPU overhead of individual torch.empty calls in the grouped linear backward passes. The existing split-quantize helpers in cast.cpp are refactored to delegate to this new primitive, and all three grouped linear backward implementations are updated to use it for weight gradient allocation.

• New allocate.cpp: Computes per-tensor byte offsets with configurable alignment, pads the base buffer for pointer alignment, and uses a shared_ptr<at::Tensor> deleter to keep the backing buffer alive as long as any view exists; zero-size tensors fall back to a standalone at::empty to avoid from_blob edge-case bugs.
• cast.cpp refactor: Removes three copies of the bespoke buffer-management lambda and delegates to bulk_allocate; the contiguous_data_and_scale pre-check logic is correctly rewritten as a per-tensor size divisibility test that is semantically equivalent to the original cumulative offset check.
• Python call-sites: Consistently pass device as the third positional argument and [256] * num_groups as alignments, matching the pybind11 signature.

Confidence Score: 5/5

The change is safe to merge: it is a pure performance optimization with well-contained scope — a new allocation helper replaces equivalent per-tensor allocations, and the refactored cast.cpp paths preserve the same memory layout semantics.

The core bulk_allocate logic is straightforward: offset arithmetic, a single CUDA allocation, and reference-counted views. The alignment padding is conservative and correct. All three Python call-sites pass arguments in the right positional order. The cast.cpp refactoring replaces identical local lambdas with equivalent calls to the new primitive, and the contiguous_data_and_scale rewrite is semantically equivalent to the original cumulative-offset check.

No files require special attention; all changes are self-contained and the critical path is covered by the three Python call-sites.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/allocate.cpp	New file implementing bulk_allocate: allocates a single contiguous CUDA buffer and creates at::from_blob views for each tensor, with per-tensor alignment support and correct empty-tensor workaround.
transformer_engine/pytorch/csrc/extensions/cast.cpp	Refactors bulk_allocate_fp8/fp4 helpers to delegate to the new bulk_allocate function; eliminates duplicated buffer-management code while preserving the contiguous_data_and_scale semantics.
transformer_engine/pytorch/csrc/extensions/pybind.cpp	Exposes bulk_allocate to Python with py::call_guardpy::gil_scoped_release(), consistent with other allocation-adjacent bindings in the file.
transformer_engine/pytorch/module/grouped_linear.py	Replaces per-expert torch.empty loop with tex.bulk_allocate for wgrad tensors, using 256-byte alignment and the context device.
transformer_engine/pytorch/ops/basic/grouped_linear.py	Replaces per-expert torch.empty loop with tex.bulk_allocate for wgrad tensors; aligned with module/grouped_linear.py change.
transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py	Replaces per-expert torch.empty loop with tex.bulk_allocate in _compute_grad_params; correctly passes device as third arg and [256]*num_groups as alignments.

Sequence Diagram

sequenceDiagram
    participant Py as Python (grouped_linear backward)
    participant PB as pybind11
    participant BA as bulk_allocate (C++)
    participant Alloc as CUDA Allocator

    Py->>PB: tex.bulk_allocate(shapes, dtypes, device, alignments)
    Note over PB: GIL released
    PB->>BA: bulk_allocate(shapes, dtypes, device, alignments)
    BA->>BA: "compute per-tensor offsets & base_alignment"
    BA->>BA: "base_byte_size += base_alignment (padding)"
    BA->>Alloc: "at::empty({base_byte_size}, kUInt8, device)"
    Alloc-->>BA: "base_buffer (shared_ptr<at::Tensor>)"
    BA->>BA: align base_ptr to base_alignment
    loop for each tensor i
        alt "byte_sizes[i] == 0"
            BA->>Alloc: at::empty(shape_i, dtype_i)
            Alloc-->>BA: standalone empty tensor
        else
            BA->>BA: at::from_blob(base_ptr+offset[i], shape_i, deleter, dtype_i)
        end
    end
    BA-->>PB: "vector<at::Tensor>"
    PB-->>Py: List[torch.Tensor] (wgrad_list)
    Note over Py: base_buffer kept alive via shared_ptr in each view deleter

Loading

_{Reviews (7): Last reviewed commit: "Merge branch 'main' into tmoon/group-mlp..." | Re-trigger Greptile}

[ptrendx]

I don't think I understand the logic here. What is the case where we would not be able to pack those buffers together if we can control the alignment requirements for the individual tensors in the allocation?

[timmoon10]

I think there was some alignment requirement for the weight tensors, and some implementation difference depending on if the weights are contiguous. In any case, this is the logic in the existing implementation and I'm trying to avoid functional changes.

[timmoon10]

/te-ci pytorch

[timmoon10]

/te-ci pytorch

[vthumbe1503]

LGTM

[timmoon10]

/te-ci pytorch