Optimize function that loads pointers on GPU

[timmoon10]

Description

tex.get_device_pointer_for_data_and_scales has two problems:

1. It has significant CPU overhead (see [PyTorch] Reduce CPU overhead in grouped MLP block #2897). In a representative benchmark on a GB200, it takes 72 us per call.
2. The meaning is extremely unintuitive. The most natural interpretation is that it takes a FP8/MXFP8/NVFP4 tensor and returns pointers as two int s. But actually it takes the buffers from multiple MXFP8/NVFP4 tensors (all assumed to have the same shape), swizzles the scaling factors, and transfers the pointers to a GPU array in a CUDA Graph-friendly way.

This PR makes several optimizations to reduce CPU overhead, mostly to reduce unnecessary heap allocations. I've also attempted to make the functionality more general and logical:

• nvte_load_value_on_device: A general function for copying a small amount of data to GPU in a CUDA Graph-friendly way. Unlike nvte_convert_pointers_to_tensor, it makes no assumptions that the data is a list of pointers.
• tex.load_data_ptrs_on_device: Takes a list of tensors and puts their data pointers into a GPU buffer.
• tex.transform_and_load_data_ptrs_on_device: Performs a user-specified transform on a list of tensors and puts the resulting data pointers into a GPU buffer. Currently it only supports scale swizzles on uniformly shaped tensors, but the transform names help make the contracts explicit.

With these changes, per-call CPU runtime has dropped from 72 us to 41 us on a GB200 node.

This is progress toward #2897.

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
• Performance optimization

Changes

• Avoid heap allocation when converting NVTEBasicTensor to transformer_engine::SimpleTensor
• Avoid heap allocation in transformer_engine::Tensor shape functions
• Add transformer_engine::Tensor::flat_2d_dims to compute first and last dims simultaneously
• Generalize and rename nvte_load_value_on_device
• Refactor and rename tex.load_data_ptrs_on_device and tex.transform_and_load_data_ptrs_on_device

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 refactors the GPU pointer-loading path for grouped MLP kernels, replacing nvte_convert_pointers_to_tensor / get_device_pointer_for_data_and_scales with a cleaner, more general API (nvte_load_value_on_device, load_data_ptrs_on_device, transform_and_load_data_ptrs_on_device) and eliminating significant heap allocation hot spots to cut per-call CPU overhead from ~72 µs to ~41 µs.

• Introduces a stack-allocated Shape class (wrapping NVTEShape) replacing std::vector<size_t> on all hot paths, and a flat_2d_dims() method that folds the two existing single-dim helpers into a single shape traversal.
• Replaces the pointer-specific CUDA kernel with a general nvte_load_value_on_device that passes data via kernel arguments for CUDA Graph compatibility, chunking at 2 KB per launch.
• Refactors the Python-side swizzle + pointer-collection helpers into two clearly-named functions and fixes the pre-existing copy-paste error in the NVFP4 columnwise_scale_inv error message.

Confidence Score: 5/5

Safe to merge; all refactored paths produce the same outputs as before with no data-corruption or functional regression identified.

The nvte_load_value_on_device kernel correctly handles aligned vector writes and byte-level tail writes with only one thread writing tail bytes. The Shape class's only resize call site explicitly sets the element after resizing, so no uninitialized reads occur today. The dropped is_cuda() guard and Shape::resize() zero-init difference are style-level observations that do not affect any currently exercised code path.

No files require special attention; the most structurally novel file is transformer_engine/common/util/utils.cu with its payload-via-kernel-args design, which is straightforward once the CUDA Graph motivation is understood.

Important Files Changed

Filename	Overview
transformer_engine/common/util/utils.cu	Replaces old pointer-specific kernel with a general nvte_load_value_on_device that passes 2 KB payloads via kernel arguments for CUDA Graph compatibility; tail-byte handling and chunking loop are correct.
transformer_engine/common/common.h	Adds stack-allocated Shape class wrapping NVTEShape, eliminating heap allocation for tensor shapes; flat_2d_dims() avoids double traversal; Shape::resize() omits zero-initialization for new elements (only call site in clear() explicitly sets the element after resize).
transformer_engine/pytorch/csrc/extensions/utils.cpp	Replaces monolithic get_device_pointer_for_data_and_scales with load_data_ptrs_on_device and transform_and_load_data_ptrs_on_device; uniform-stride buffer avoids per-tensor offset tracking; dropped is_cuda() guard noted.
transformer_engine/common/transformer_engine.cpp	Shape-validation logic refactored to std::array-based; NVFP4 columnwise error message now correctly references t.columnwise_scale_inv.shape.
transformer_engine/common/swizzle/swizzle.cu	flat_2d_dims() consolidation and reserve() additions are correct micro-optimizations.
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py	Callers correctly split; _fc*_sfb_buffer keepalive pattern is correct for CPython reference-counting.
transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py	Same caller migration as forward path; keepalive pattern is correct.

Sequence Diagram

sequenceDiagram
    participant PY as Python caller
    participant LD as load_data_ptrs_on_device
    participant TR as transform_and_load_data_ptrs_on_device
    participant NV as nvte_load_value_on_device
    participant SW as nvte_multi_tensor_swizzle_scaling_factors
    participant GPU as Device memory

    PY->>LD: tensors, device
    LD->>LD: collect data_ptr() to ptrs_host[]
    LD->>NV: "ptrs_host, ptrs_device, N*8 bytes"
    NV->>GPU: kernel (payload in args) writes pointers
    LD-->>PY: ptrs_device tensor

    PY->>TR: transform_type, scale tensors, device
    TR->>TR: allocate swizzled_scales buffer (uniform stride)
    TR->>SW: input/output TensorWrapper pairs
    SW->>GPU: swizzle kernel writes swizzled scales
    TR->>TR: collect swizzled_scales_dptr offsets to ptrs_host[]
    TR->>NV: "ptrs_host, ptrs_device, N*8 bytes"
    NV->>GPU: kernel (payload in args) writes pointers
    TR-->>PY: (ptrs_device, swizzled_scales keepalive)

Loading

_{Reviews (5): Last reviewed commit: "Make Shape conversion constructors expli..." | Re-trigger Greptile}

[timmoon10]

Without this assignment operator, assigning from a NVTEBasicTensor triggers a heap allocator in the NVTEBasicTensor constructor. We do this assignment frequently within nvte_set_tensor_param_v2.

[timmoon10]

I'm not committed to this name. I based it on std::transform. I suppose "map" would be more Python-focused, but that sounds worse.

[timmoon10]

/te-ci

[ptrendx]

Seems a lot of those changes would basically not be needed if we did not use the std::vector in Tensor/SimpleTensor and just used NVTEShape everywhere - this would effectively make SimpleTensor and NVTEBasicTensor the same thing (we could even do the constructor in the public header, just behing the if cplusplus guard).

[vthumbe1503]

Thanks for cleaning up the APIs. Looks much nicer now. CPU overheads being caused by heap allocations of shape, makes me wonder whether we should revive this PR to standardize on NVTEShape yo avoid back and forth between vector<size_t> and NVTE_Shape

[vthumbe1503]

The other optimization can be to load both fc1 and fc2 data and scale inv togegther at the start of backward. I am hoping it wouldnt make the code ugly.

[vthumbe1503]

So when you say heap allocations being done redundantly again and again. Do you mean the vector to NVTEShape conversions?

I rememember this problem being observed even with a basic te linear profiling. And I hadnt gotten this PR merged.
#2514

which essentially standadizes to use NVTEShape everywhere instead of using vector at all to avoid bouncing back and forth between the two allocations. Maybe it might be worth to revive the PR?

[timmoon10]

Previously, assigning an NVTEBasicTensor to a SimpleTensor would trigger the constructor and then the move operator. This would allocate an std::vector, move it, and deallocate the old std::vector.

One other approach I was thinking about was implementing a Shape class that wraps around NVTEShape and has a similar API as std::vector. That way we can keep the nice ergonomics, while avoiding heap allocations.

[vthumbe1503]

I had tried your other approach in the 2514 PR above, but eventually had removed it due to some complications. I have refer back to my notes on why it didnt work out for me.

But here is the commit that reverted it
b599776

I had called it NVTEShapeWrapper and implemented all the vector based APIs.

[vthumbe1503]

One complication I do remember was to change a lot of attention interfaces to have NVTEShapeWrapper instead of using vector.

[timmoon10]

I've found that adding a cast operator to std::vector helps reduce the number of places we need to change the interfaces.

[timmoon10]

/te-ci