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>

Author: timmoon10

transformer_engine/common/util/utils.cu

@@ -19,21 +19,22 @@ namespace load_value_on_device {
 namespace {
 
 union Payload {
-  static constexpr size_t max_bytes = 2048;
-  static constexpr size_t vector_size = 4;
-  uint8_t bytes[max_bytes];
-  uint32_t vectors[max_bytes / vector_size];
+  static constexpr size_t kMaxBytes = 2048;
+  static constexpr size_t kVectorSize = 4;
+  static constexpr size_t kMaxVectors = kMaxBytes / kVectorSize;
+  uint8_t bytes[kMaxBytes];
+  uint32_t vectors[kMaxVectors];
 };
 
 constexpr size_t block_size = 512;
-constexpr size_t num_blocks = DIVUP(Payload::max_bytes / Payload::vector_size, block_size);
+constexpr size_t num_blocks = DIVUP(Payload::kMaxBytes / Payload::kVectorSize, block_size);
 
 __global__ void __launch_bounds__(block_size) kernel(Payload payload, size_t num_bytes, void *out) {
   const size_t tid = blockIdx.x * blockDim.x + threadIdx.x;
-  if (Payload::vector_size * (tid + 1) <= num_bytes) {
+  if (Payload::kVectorSize * (tid + 1) <= num_bytes) {
     reinterpret_cast<uint32_t *>(out)[tid] = payload.vectors[tid];
   } else {
-    for (int64_t i = Payload::vector_size * tid; i < num_bytes; ++i) {
+    for (size_t i = Payload::kVectorSize * tid; i < num_bytes; ++i) {
       static_cast<uint8_t *>(out)[i] = payload.bytes[i];
     }
   }
@@ -56,15 +57,15 @@ void nvte_load_value_on_device(const void *host_ptr, void *device_ptr, size_t nu
   // Check pointers
   NVTE_CHECK(host_ptr != nullptr, "Attempting to read ", num_bytes, " bytes from a null pointer.");
   NVTE_CHECK(device_ptr != nullptr, "Attempting to write ", num_bytes, " bytes into a null pointer.");
-  NVTE_CHECK(reinterpret_cast<uintptr_t>(device_ptr) % Payload::vector_size == 0,
-             "Device pointer is not aligned to ", Payload::vector_size, " bytes.");
+  NVTE_CHECK(reinterpret_cast<uintptr_t>(device_ptr) % Payload::kVectorSize == 0,
+             "Device pointer is not aligned to ", Payload::kVectorSize, " bytes.");
 
   // Chunk data to fit in kernel arguments and launch kernels
   const uint8_t *src = static_cast<const uint8_t *>(host_ptr);
   uint8_t *dst = static_cast<uint8_t *>(device_ptr);
-  for (size_t offset = 0; offset < num_bytes; offset += Payload::max_bytes) {
-    const size_t chunk_size = std::min(num_bytes - offset, Payload::max_bytes);
-    Payload payload;
+  for (size_t offset = 0; offset < num_bytes; offset += Payload::kMaxBytes) {
+    const size_t chunk_size = std::min(num_bytes - offset, Payload::kMaxBytes);
+    Payload payload{};
     std::memcpy(payload.bytes, src + offset, chunk_size);
     kernel<<<num_blocks, block_size, 0, stream>>>(payload, chunk_size, dst + offset);
     NVTE_CHECK_CUDA(cudaGetLastError());

transformer_engine/pytorch/csrc/extensions/utils.cpp

@@ -27,7 +27,7 @@ at::Tensor load_data_ptrs_on_device(const std::vector<at::Tensor> &tensors,
   }
 
   // Allocate device buffer
-  auto ptrs_device = at::empty({tensors.size()},
+  auto ptrs_device = at::empty({static_cast<int64_t>(tensors.size())},
                                at::TensorOptions().dtype(at::kLong).device(device));
 
   // Load pointers on device
@@ -51,7 +51,7 @@ std::tuple<at::Tensor, std::optional<at::Tensor>> transform_and_load_data_ptrs_o
   if (num_tensors == 0) {
     // No input tensors, return tensor with no elements
     return {
-      at::empty(std::vector<int64_t>{0}, at::TensorOptions().dtype(at::kLong).device(device)),
+      at::empty({int64_t{0}}, at::TensorOptions().dtype(at::kLong).device(device)),
       std::nullopt};
   }
 
@@ -114,10 +114,10 @@ std::tuple<at::Tensor, std::optional<at::Tensor>> transform_and_load_data_ptrs_o
       NVTE_ERROR("Unsupported case.");
     }
 
-    // Allocate single buffer for swizzled scales
+    // Allocate single buffer for swizzled scales.
+    // Uses a uniform stride since all tensors share the same scale shape.
     const size_t swizzled_scales_stride = roundup(scale_bytes, 16);  // Align to 16 bytes
-    const int64_t swizzled_scales_bytes = swizzled_scales_stride * num_tensors;
-    auto swizzled_scales = at::empty({swizzled_scales_bytes},
+    auto swizzled_scales = at::empty({static_cast<int64_t>(swizzled_scales_stride * num_tensors)},
                                      at::TensorOptions().dtype(at::kByte).device(device));
     uint8_t *swizzled_scales_dptr = reinterpret_cast<uint8_t *>(swizzled_scales.data_ptr());
 
@@ -167,7 +167,7 @@ std::tuple<at::Tensor, std::optional<at::Tensor>> transform_and_load_data_ptrs_o
     }
 
     // Load pointers on device
-    auto ptrs_device = at::empty(std::vector<int64_t>{num_tensors},
+    auto ptrs_device = at::empty({static_cast<int64_t>(num_tensors)},
                                  at::TensorOptions().dtype(at::kLong).device(device));
     nvte_load_value_on_device(ptrs_host.data(), ptrs_device.data_ptr(),
                               num_tensors * sizeof(uint64_t),

transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py

@@ -502,7 +502,7 @@ def fuser_backward(
                 [w._columnwise_data for w in grouped_fc2_weight],
                 device,
             )
-            fc2_sfb_ptrs, fc2_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
+            fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
                 "uniform_mxfp8_columnwise_swizzle",
                 [w._columnwise_scale_inv for w in grouped_fc2_weight],
                 device,
@@ -661,7 +661,7 @@ def fuser_backward(
                     [w._columnwise_data for w in grouped_fc1_weight],
                     device,
                 )
-                fc1_sfb_ptrs, fc1_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
+                fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
                     "uniform_mxfp8_columnwise_swizzle",
                     [w._columnwise_scale_inv for w in grouped_fc1_weight],
                     device,

transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py

@@ -337,7 +337,7 @@ def fuser_forward(
                 [w._rowwise_data for w in grouped_fc1_weight],
                 device,
             )
-            fc1_sfb_ptrs, fc1_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
+            fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
                 "uniform_mxfp8_rowwise_swizzle",
                 [w._rowwise_scale_inv for w in grouped_fc1_weight],
                 device,
@@ -438,7 +438,7 @@ def fuser_forward(
                 [w._rowwise_data for w in grouped_fc2_weight],
                 device,
             )
-            fc2_sfb_ptrs, fc2_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
+            fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_load_data_ptrs_on_device(
                 "uniform_mxfp8_rowwise_swizzle",
                 [w._rowwise_scale_inv for w in grouped_fc2_weight],
                 device,