Optimize function that loads pointers on GPU (#3001)

* Remove unnecessary heap allocations

Avoid constructing temporary std::vector when converting NVTEBasicTensor to SimpleTensor. Avoid string operations in multi-tensor swizzle. Avoid temporary std::vector when checking scale tensors.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Avoid heap allocation in Tensor::flat_first_dim/flat_last_dim

Tensor::shape() returns a std::vector<size_t> by value, allocating
on the heap. flat_first_dim and flat_last_dim only need to walk
the dims, so the allocation was pure overhead in hot paths.

Introduce Tensor::compute_shape() returning an NVTEShape (fixed
inline buffer, no heap) as the single source of truth for the
format-dependent shape logic. shape() is now a thin std::vector
wrapper around it for callers that want a vector; flat_first_dim
and flat_last_dim call compute_shape() directly.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add Tensor::flat_2d_dims() to compute both matrix dims in one pass

flat_first_dim() and flat_last_dim() each called compute_shape()
independently. flat_2d_dims() computes both in a single pass; the
scalar helpers now delegate to it.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use flat_2d_dims() throughout common lib

Replace all paired flat_first_dim() + flat_last_dim() calls on the
same tensor with a single flat_2d_dims() call. Saves one compute_shape()
per tensor in CheckScaleTensorShape, the multi-tensor swizzle loop, and
various cast/GEMM dispatch paths.

Also adds reserve() to the local vectors in
nvte_multi_tensor_swizzle_scaling_factors to avoid reallocation.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Generalize API for CUDA-Graph-safe copy to GPU.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Dedup swizzle logic in get_device_pointer_for_data_and_scales

Replace the inline swizzle implementation with a call to
multi_tensor_swizzle_scales_for_gemm, which has identical logic
(16B-aligned contiguous output buffer, TensorWrapper construction,
nvte_multi_tensor_swizzle_scaling_factors kernel). Swizzled pointers
are read back from the updated TensorWrappers after the call.

Add reserve() to vectors in multi_tensor_swizzle_scales_for_gemm_impl
now that this function is on the hot path for get_device_pointer_for_data_and_scales.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Make separate functions for load data_ptrs and swizzle + load data_ptrs.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Change function name to nvte_load_value_on_device

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Fix code review issues before opening PR

- Use size_t in kernel tail loop (was int64_t)
- Zero-initialize Payload before memcpy (Payload{})
- Rename Payload members to kMaxBytes/kVectorSize/kMaxVectors (linter)
- Consistent at::empty shape pattern: {static_cast<int64_t>(N)}
- Drop intermediate swizzled_scales_bytes variable
- Add comment explaining uniform-stride assumption in
  transform_and_load_data_ptrs_on_device
- Rename sfb_buffer -> _sfb_buffer (keepalive, not directly used)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Formatter and review suggestions from @greptile-apps

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add Shape class wrapping NVTEShape

Provides a std::vector<size_t>-like interface around NVTEShape without
heap allocation, used as the return type of Tensor::shape() in place of
the previous std::vector. Disambiguate cute::Shape from
transformer_engine::Shape in the hadamard_transform kernels.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Make SimpleTensor stack-allocatable

Store shape in Shape class rather than std::vector.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Make Shape conversion constructors explicit

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Make conversion from Shape to std::vector explicit

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add batched NVTETensor create/destroy

Expose nvte_create_tensors and nvte_destroy_tensors so multi-tensor
callers can amortize the TensorAllocator mutex across N tensors
instead of locking once per call. nvte_destroy_tensors was already
defined internally but not declared in the public header.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use batched NVTETensor allocator in transform_and_load_data_ptrs_on_device

The uniform swizzle path constructed 2N TensorWrappers and then
extracted their raw NVTETensors into separate vectors. Replace with a
single 2N nvte_create_tensors call into one contiguous buffer (inputs
in the first half, outputs in the second), an RAII guard for
nvte_destroy_tensors, and a local set_param lambda for the setters.
Drops the separate pack pass and reduces the allocator mutex
acquisitions from 4N to 2 per call.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Expand usage of batched NVTETensor allocator

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Use string_view in tensor checking functions

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Tweak function names

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Pass std::string_view by value in Check*Tensor helpers

string_view is already a (ptr, len) reference — passing by const-ref
adds an indirection without benefit. Matches the C++ Core Guidelines
F.16 recommendation.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Review suggestions from @ptrendx

Expand internal usage of Shape class. Zero-initialize in Shape::resize. Make sure dynamic smem querying is per-device. Reuse logic for batched and single tensor alloc/dealloc.

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add MultiTensorWrapper for batched NVTETensor allocation

Thin RAII wrapper around a batched nvte_create_tensors /
nvte_destroy_tensors pair, with operator[], data(), iteration, and
implicit conversion to NVTETensor* for multi-tensor C APIs. Replaces
the ad-hoc DestroyGuard struct used at each call site in
recipe.cpp, swizzle.cpp, and utils.cpp.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

transformer_engine/common/cast/dispatch/quantize.cuh

@@ -98,8 +98,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
       CheckOutputTensor(*output_tensor, "output", false);
 
       // Choose kernel
-      int32_t rows = input_tensor->flat_first_dim();
-      int32_t cols = input_tensor->flat_last_dim();
+      const auto [rows, cols] = input_tensor->flat_2d_dims();
       auto dtype = input_tensor->dtype();
       const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4;
       const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled;
@@ -260,8 +259,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
       CheckOutputTensor(*output_tensor, "output", false);
 
       // Choose kernel
-      int32_t rows = grad_tensor->flat_first_dim();
-      int32_t cols = grad_tensor->flat_last_dim();
+      const auto [rows, cols] = grad_tensor->flat_2d_dims();
       auto dtype = grad_tensor->dtype();
       const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled;
       NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448,
@@ -396,8 +394,7 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou
       // output list here is allowed to have empty tensor
 
       // Choose kernel
-      int32_t rows = input_tensor->flat_first_dim();
-      int32_t cols = input_tensor->flat_last_dim();
+      const auto [rows, cols] = input_tensor->flat_2d_dims();
       auto dtype = input_tensor->dtype();
 
       const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled;

transformer_engine/common/cast/fp8/quantize_fp8.cuh

@@ -391,8 +391,7 @@ void quantize_2D(const Tensor &input, const Tensor *act_input, Tensor *output, T
   using namespace quantize_2D_kernel;
   checkCuDriverContext(stream);
 
-  const size_t rows = input.flat_first_dim();
-  const size_t cols = input.flat_last_dim();
+  const auto [rows, cols] = input.flat_2d_dims();
   const size_t chunks_Y = DIVUP(rows, FP8_CHUNK_DIM_Y);
   const size_t chunks_X = DIVUP(cols, FP8_CHUNK_DIM_X);
   const size_t blocks_Y = chunks_Y;
@@ -406,7 +405,7 @@ void quantize_2D(const Tensor &input, const Tensor *act_input, Tensor *output, T
 
   if constexpr (IS_DBIAS) {
     NVTE_CHECK(dbias->data.dtype == input.data.dtype, "DBias must have the same type as input.");
-    NVTE_CHECK(dbias->data.shape == std::vector<size_t>{cols}, "Wrong shape of DBias.");
+    NVTE_CHECK(dbias->data.shape == Shape{cols}, "Wrong shape of DBias.");
     NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor.");
 
     if (workspace->data.dptr == nullptr) {

transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh

@@ -261,8 +261,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream)
   const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1;
   const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1;
 
-  const size_t rows = input.flat_first_dim();
-  const size_t cols = input.flat_last_dim();
+  const auto [rows, cols] = input.flat_2d_dims();
   const size_t chunks_Y = DIVUP(rows, CHUNK_DIM_Y);
   const size_t chunks_X = DIVUP(cols, CHUNK_DIM_X);
 

transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh

@@ -867,7 +867,7 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations
     NVTE_CHECK(dbias->data.dtype == input->dtype(),
                "DBias must have the same type as input_tensor.");
 
-    std::vector<size_t> expected_shape_dbias_tensor = {num_tensors, last_logical_dim};
+    Shape expected_shape_dbias_tensor = {num_tensors, last_logical_dim};
     NVTE_CHECK(dbias->data.shape == expected_shape_dbias_tensor, "Wrong shape of DBias.");
 
     NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor.");

transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh

@@ -578,8 +578,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
   constexpr bool CAST_DBIAS_ONLY = IS_DBIAS && (!IS_DACT) && (!IS_ACT);
 
   // Tensor dimensions
-  const size_t rows = input.flat_first_dim();
-  const size_t cols = input.flat_last_dim();
+  const auto [rows, cols] = input.flat_2d_dims();
 
   // Tensor chunk handled by each CUDA block
   constexpr size_t CHUNK_DIM_Y = CAST_DBIAS_ONLY ? 128 : 64;
@@ -622,7 +621,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
 
   if constexpr (IS_DBIAS) {
     NVTE_CHECK(dbias->data.dtype == input.dtype(), "DBias must have the same type as input.");
-    NVTE_CHECK(dbias->data.shape == std::vector<size_t>{cols}, "Wrong shape of DBias.");
+    NVTE_CHECK(dbias->data.shape == Shape{cols}, "Wrong shape of DBias.");
     NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor.");
 
     if (workspace->data.dptr == nullptr) {

transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh

@@ -95,8 +95,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream)
   const int e4m3_max = input.nvfp4_e4m3_max;
 
   constexpr int FP4_BLOCK_SIZE = 16;
-  const size_t N = input.flat_first_dim();
-  const size_t M = input.flat_last_dim();
+  const auto [N, M] = input.flat_2d_dims();
 
   NVTE_CHECK(M % FP4_BLOCK_SIZE == 0, "Last dimension of FP4 tensors needs to be divisible by ",
              FP4_BLOCK_SIZE, ", but got ", input.data.shape, ".");

transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh

@@ -783,8 +783,7 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop,
 
   NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data.");
 
-  const size_t rows = input.flat_first_dim();
-  const size_t cols = input.flat_last_dim();
+  const auto [rows, cols] = input.flat_2d_dims();
 
   NVTE_CHECK(rows % 32 == 0,
              "Number of tensor rows must be a multiple of 32");  // 16B alignment for TMA
@@ -835,7 +834,7 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop,
     Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor);
     NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64,
                "RNG state should contain 2 64-bit values.");
-    NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector<size_t>{2},
+    NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2},
                "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape);
     rng_state = reinterpret_cast<const size_t *>(rng_state_te_tensor.data.dptr);
   }

transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh

@@ -121,8 +121,7 @@ inline void compute_rowwise_amax(const Tensor &input, const Tensor *noop, Tensor
 #if FP4_TYPE_SUPPORTED
   using namespace rowwise_amax_kernel;
 
-  const size_t rows = input.flat_first_dim();
-  const size_t cols = input.flat_last_dim();
+  const auto [rows, cols] = input.flat_2d_dims();
   NVTE_CHECK(cols % ROWWISE_AMAX_SF_VEC_SIZE == 0,
              "Row-scaled NVFP4 quantization requires last dim divisible by ",
              ROWWISE_AMAX_SF_VEC_SIZE, ".");
@@ -1359,8 +1358,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output,
                "Transposed scaling tensor must be allocated");
   }
 
-  const size_t rows = input.flat_first_dim();
-  const size_t cols = input.flat_last_dim();
+  const auto [rows, cols] = input.flat_2d_dims();
 
   NVTE_CHECK(rows % 32 == 0,
              "Number of tensor rows must be a multiple of 32");  // 16B alignment for TMA
@@ -1391,7 +1389,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output,
     Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor);
     NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64,
                "RNG state should contain 2 64-bit values.");
-    NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector<size_t>{2},
+    NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2},
                "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape);
     rng_state = reinterpret_cast<const size_t *>(rng_state_te_tensor.data.dptr);
   }

transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh

@@ -718,8 +718,7 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop,
                "Transposed scaling tensor must be allocated");
   }
 
-  const size_t rows = input.flat_first_dim();
-  const size_t cols = input.flat_last_dim();
+  const auto [rows, cols] = input.flat_2d_dims();
 
   NVTE_CHECK(rows % 32 == 0,
              "Number of tensor rows must be a multiple of 32");  // 16B alignment for TMA
@@ -750,7 +749,7 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop,
     Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor);
     NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64,
                "RNG state should contain 2 64-bit values.");
-    NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector<size_t>{2},
+    NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2},
                "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape);
     rng_state = reinterpret_cast<const size_t *>(rng_state_te_tensor.data.dptr);
   }

transformer_engine/common/comm_gemm/comm_gemm.cpp

@@ -130,12 +130,9 @@ int64_t block_size(NVTECommGemmCtx* ctx, int64_t global_size) {
 void AgGemmInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n, int64_t k,
                         const Tensor* a, const Tensor* b, const Tensor* d, bool transa,
                         bool transb) {
-  const auto a0 = a->flat_first_dim();
-  const auto a1 = a->flat_last_dim();
-  const auto b0 = b->flat_first_dim();
-  const auto b1 = b->flat_last_dim();
-  const auto d0 = d->flat_first_dim();
-  const auto d1 = d->flat_last_dim();
+  const auto [a0, a1] = a->flat_2d_dims();
+  const auto [b0, b1] = b->flat_2d_dims();
+  const auto [d0, d1] = d->flat_2d_dims();
 
   if (transa) {
     NVTE_CHECK(a1 == k, "Unsupported tensor dimension in A: expected ", k, ", got ", a1);
@@ -169,12 +166,9 @@ void AgGemmInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n
 void GemmRsInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n, int64_t k,
                         const Tensor* a, const Tensor* b, const Tensor* d, bool transa,
                         bool transb) {
-  const auto a0 = a->flat_first_dim();
-  const auto a1 = a->flat_last_dim();
-  const auto b0 = b->flat_first_dim();
-  const auto b1 = b->flat_last_dim();
-  const auto d0 = d->flat_first_dim();
-  const auto d1 = d->flat_last_dim();
+  const auto [a0, a1] = a->flat_2d_dims();
+  const auto [b0, b1] = b->flat_2d_dims();
+  const auto [d0, d1] = d->flat_2d_dims();
 
   if (transa) {
     NVTE_CHECK(a0 == m, "Unsupported tensor dimension in A: expected ", m, ", got ", a0);
@@ -213,12 +207,9 @@ void GemmRsInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n
 void GemmArInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n, int64_t k,
                         const Tensor* a, const Tensor* b, const Tensor* d, bool transa,
                         bool transb) {
-  const auto a0 = a->flat_first_dim();
-  const auto a1 = a->flat_last_dim();
-  const auto b0 = b->flat_first_dim();
-  const auto b1 = b->flat_last_dim();
-  const auto d0 = d->flat_first_dim();
-  const auto d1 = d->flat_last_dim();
+  const auto [a0, a1] = a->flat_2d_dims();
+  const auto [b0, b1] = b->flat_2d_dims();
+  const auto [d0, d1] = d->flat_2d_dims();
 
   if (transa) {
     NVTE_CHECK(a0 == m, "Unsupported tensor dimension in A: expected ", m, ", got ", a0);