[PyTorch]Validate grouped split alignment on device

transformer_engine/common/include/transformer_engine/transformer_engine.h

@@ -546,6 +546,19 @@ void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs,
                                   const int64_t *strides, const int *include_leading_zero,
                                   size_t num_outputs, cudaStream_t stream);
 
+/*! \brief Compute multiple scaled offsets and assert row split alignment on the device.
+ *
+ * This uses the same prefix-sum kernel as nvte_splits_to_offsets_multi and does not synchronize
+ * the host, so the validation remains active during CUDA graph replay.
+ *
+ * \param[in] split_alignment Positive required divisor of every split size.
+ * \see nvte_splits_to_offsets_multi for the remaining parameters.
+ */
+void nvte_splits_to_offsets_multi_v2(NVTETensor split_sizes, size_t split_alignment,
+                                     NVTETensor *outputs, const int64_t *strides,
+                                     const int *include_leading_zero, size_t num_outputs,
+                                     cudaStream_t stream);
+
 /*! \brief TE Grouped Tensor type
  *
  * NVTEGroupedTensor is a collection of tensors with potentially different shapes

transformer_engine/common/util/splits_to_offsets.cu

@@ -18,6 +18,7 @@ namespace transformer_engine::splits_to_offsets {
 namespace {
 
 constexpr size_t kThreadsPerBlock = 256;
+constexpr size_t kSmallThreadsPerBlock = 128;
 
 struct KernelArgs {
   static constexpr size_t kMaxNumOutputs = 8;
@@ -29,6 +30,7 @@ struct KernelArgs {
   int64_t strides[kMaxNumOutputs] = {};
   bool include_leading_zero[kMaxNumOutputs] = {};
   size_t num_outputs = 0;
+  size_t split_alignment = 1;
 };
 
 __device__ __forceinline__ int64_t load_split_size(const void *__restrict__ split_sizes,
@@ -58,7 +60,8 @@ __device__ __forceinline__ void store_output(void *__restrict__ output, DType dt
   }
 }
 
-__global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) {
+template <size_t kScanThreads>
+__global__ void __launch_bounds__(kScanThreads) kernel(KernelArgs args) {
   const size_t tid = threadIdx.x;
 
   // Fill leading zeros if needed
@@ -71,7 +74,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) {
   }
 
   // Workspace for prefix sum chunk
-  __shared__ int64_t block_scan[kThreadsPerBlock];
+  __shared__ int64_t block_scan[kScanThreads];
 
   // Sum from previous chunks
   __shared__ int64_t chunk_prefix;
@@ -81,19 +84,35 @@ __global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) {
   __syncthreads();
 
   // Perform prefix sum in chunks
-  for (size_t chunk_start = 0; chunk_start < args.num_splits; chunk_start += kThreadsPerBlock) {
+  for (size_t chunk_start = 0; chunk_start < args.num_splits; chunk_start += kScanThreads) {
     const size_t idx = chunk_start + tid;
 
-    // Load input from global memory into shared memory
-    if (idx < args.num_splits) {
-      block_scan[tid] = load_split_size(args.split_sizes, args.split_sizes_dtype, idx);
-    } else {
-      block_scan[tid] = 0;
+    // Validate before the scan. All threads leave together so no thread waits at a barrier
+    // after another thread has trapped on invalid split metadata.
+    const int64_t split =
+        idx < args.num_splits ? load_split_size(args.split_sizes, args.split_sizes_dtype, idx) : 0;
+    if (args.split_alignment > 1) {
+      const size_t alignment_mask = args.split_alignment - 1;
+      bool misaligned;
+      if ((args.split_alignment & alignment_mask) == 0) {
+        misaligned = (static_cast<size_t>(split) & alignment_mask) != 0;
+      } else {
+        misaligned = static_cast<size_t>(split) % args.split_alignment != 0;
+      }
+      const bool invalid = idx < args.num_splits && (split < 0 || misaligned);
+      if (__syncthreads_or(invalid)) {
+        if (tid == 0) {
+          NVTE_DEVICE_ERROR("Grouped split size is not divisible by split_alignment.");
+          __trap();  // NVTE_DEVICE_ERROR's assert is compiled out of release builds.
+        }
+        return;
+      }
     }
+    block_scan[tid] = split;
     __syncthreads();
 
     // Prefix sum in shared memory
-    for (size_t offset = 1; offset < kThreadsPerBlock; offset <<= 1) {
+    for (size_t offset = 1; offset < kScanThreads; offset <<= 1) {
       const int64_t addend = (tid >= offset) ? block_scan[tid - offset] : 0;
       __syncthreads();
       block_scan[tid] += addend;
@@ -112,7 +131,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) kernel(KernelArgs args) {
 
     // Update sum for later chunks
     __syncthreads();
-    if (tid == kThreadsPerBlock - 1) {
+    if (tid == kScanThreads - 1) {
       chunk_prefix += block_scan[tid];
     }
     __syncthreads();
@@ -187,7 +206,11 @@ void nvte_splits_to_offsets(const int64_t *split_sizes, int64_t *output, size_t
   args.strides[0] = stride;
   args.include_leading_zero[0] = true;
   args.num_outputs = 1;
-  s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args);
+  if (num_splits <= s2o::kSmallThreadsPerBlock) {
+    s2o::kernel<s2o::kSmallThreadsPerBlock><<<1, s2o::kSmallThreadsPerBlock, 0, stream>>>(args);
+  } else {
+    s2o::kernel<s2o::kThreadsPerBlock><<<1, s2o::kThreadsPerBlock, 0, stream>>>(args);
+  }
   NVTE_CHECK_CUDA(cudaGetLastError());
 }
 
@@ -230,10 +253,19 @@ void nvte_splits_to_offsets_2d(const NVTETensor first_dims, const NVTETensor las
 void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs,
                                   const int64_t *strides, const int *include_leading_zero,
                                   size_t num_outputs, cudaStream_t stream) {
-  NVTE_API_CALL(nvte_splits_to_offsets_multi);
+  nvte_splits_to_offsets_multi_v2(split_sizes, 1, outputs, strides, include_leading_zero,
+                                  num_outputs, stream);
+}
+
+void nvte_splits_to_offsets_multi_v2(NVTETensor split_sizes, size_t split_alignment,
+                                     NVTETensor *outputs, const int64_t *strides,
+                                     const int *include_leading_zero, size_t num_outputs,
+                                     cudaStream_t stream) {
+  NVTE_API_CALL(nvte_splits_to_offsets_multi_v2);
   using namespace transformer_engine;
   namespace s2o = transformer_engine::splits_to_offsets;
 
+  NVTE_CHECK(split_alignment > 0, "split_alignment must be positive.");
   if (num_outputs == 0) {
     return;
   }
@@ -279,14 +311,19 @@ void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs,
     args.split_sizes_dtype = split_sizes_dtype;
     args.num_splits = num_splits;
     args.num_outputs = chunk_size;
+    args.split_alignment = split_alignment;
     for (size_t i = 0; i < chunk_size; ++i) {
       const size_t out_idx = chunk_start + i;
       args.outputs[i] = output_tensors[out_idx]->data.dptr;
       args.outputs_dtype[i] = output_tensors[out_idx]->dtype();
       args.strides[i] = strides[out_idx];
       args.include_leading_zero[i] = include_leading_zero[out_idx] != 0;
     }
-    s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args);
+    if (num_splits <= s2o::kSmallThreadsPerBlock) {
+      s2o::kernel<s2o::kSmallThreadsPerBlock><<<1, s2o::kSmallThreadsPerBlock, 0, stream>>>(args);
+    } else {
+      s2o::kernel<s2o::kThreadsPerBlock><<<1, s2o::kThreadsPerBlock, 0, stream>>>(args);
+    }
     NVTE_CHECK_CUDA(cudaGetLastError());
   }
 }

transformer_engine/pytorch/csrc/extensions.h

@@ -586,7 +586,7 @@ at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_
 std::tuple<at::Tensor, std::vector<at::Tensor>> splits_to_offsets_multi(
     const at::Tensor &split_sizes, const c10::Device &device, const std::vector<int64_t> &strides,
     const std::vector<bool> &include_leading_zero, const std::vector<at::ScalarType> &dtypes,
-    bool bulk_allocate_outputs);
+    bool bulk_allocate_outputs, size_t split_alignment);
 
 at::Tensor copy_data_ptrs_to_device(const std::vector<at::Tensor> &tensors,
                                     const c10::Device &device);

transformer_engine/pytorch/csrc/extensions/misc.cpp

@@ -42,7 +42,7 @@ at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_
 std::tuple<at::Tensor, std::vector<at::Tensor>> splits_to_offsets_multi(
     const at::Tensor &split_sizes, const c10::Device &device, const std::vector<int64_t> &strides,
     const std::vector<bool> &include_leading_zero, const std::vector<at::ScalarType> &dtypes,
-    bool bulk_allocate_outputs) {
+    bool bulk_allocate_outputs, size_t split_alignment) {
   const size_t num_outputs = strides.size();
   const size_t num_splits = static_cast<size_t>(split_sizes.numel());
 
@@ -51,6 +51,7 @@ std::tuple<at::Tensor, std::vector<at::Tensor>> splits_to_offsets_multi(
              "strides, include_leading_zero, and dtypes must have matching lengths, but got ",
              strides.size(), ", ", include_leading_zero.size(), ", and ", dtypes.size(), ".");
   NVTE_CHECK(device.is_cuda(), "device must be CUDA, but got ", device.str(), ".");
+  NVTE_CHECK(split_alignment > 0, "split_alignment must be positive.");
 
   // Convert split sizes to int64 GPU tensor.
   const at::Tensor split_sizes_i64 =
@@ -94,9 +95,9 @@ std::tuple<at::Tensor, std::vector<at::Tensor>> splits_to_offsets_multi(
 
   auto split_sizes_nvte = makeTransformerEngineTensor(split_sizes_out);
   NVTE_SCOPED_GIL_RELEASE({
-    nvte_splits_to_offsets_multi(split_sizes_nvte.data(), outputs_nvte.data(), strides.data(),
-                                 include_leading_zero_int.data(), num_outputs,
-                                 at::cuda::getCurrentCUDAStream());
+    nvte_splits_to_offsets_multi_v2(split_sizes_nvte.data(), split_alignment, outputs_nvte.data(),
+                                    strides.data(), include_leading_zero_int.data(), num_outputs,
+                                    at::cuda::getCurrentCUDAStream());
   });
 
   return {split_sizes_out, std::move(outputs)};

transformer_engine/pytorch/csrc/extensions/pybind.cpp

@@ -625,7 +625,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("splits_to_offsets_multi", &transformer_engine::pytorch::splits_to_offsets_multi,
         "Compute multiple scaled inclusive-scan offsets from a split-sizes vector",
         py::arg("split_sizes"), py::arg("device"), py::kw_only(), py::arg("strides"),
-        py::arg("include_leading_zero"), py::arg("dtypes"), py::arg("bulk_allocate") = false);
+        py::arg("include_leading_zero"), py::arg("dtypes"), py::arg("bulk_allocate") = false,
+        py::arg("split_alignment") = 1);
   m.def("get_num_cublas_streams", &nvte_get_num_compute_streams, "Get number of compute streams",
         py::call_guard<py::gil_scoped_release>());
 

transformer_engine/pytorch/ops/basic/grouped_linear.py

@@ -155,8 +155,7 @@ class GroupedLinear(BasicOperation):
     The grouped-tensor path is selected automatically where it is
     supported (see ``is_op_fuser_grouped_tensor_path_supported``). With
     MXFP8 or FP8 block scaling, it relies on every split size being a
-    multiple of 128, as ``split_alignment`` declares. This is checked only
-    on the device, so other split sizes can produce incorrect results.
+    multiple of 128, as ``split_alignment`` declares.
 
     Parameters
     ----------
@@ -201,7 +200,8 @@ class GroupedLinear(BasicOperation):
         in the forward pass. The scale tensor has shape
         ``(total_tokens,)`` and is split according to the split sizes.
     split_alignment : int, default = ``128``
-        Guaranteed alignment of each input split size.
+        Expected alignment for input split sizes. Split sizes that do not
+        match alignment may produce errors or incorrect results.
     """
 
     # Operation expects input split sizes (and optionally scales tensor)
@@ -1363,6 +1363,7 @@ def _fuser_forward_grouped_tensor(
             include_leading_zero=[False, True, True, True],
             dtypes=[torch.int32, torch.int64, torch.int64, torch.int64],
             bulk_allocate=True,
+            split_alignment=self.split_alignment,
         )
         split_points = grouped_tensor_offsets[0]
         base_split_offsets = grouped_tensor_offsets[1]

transformer_engine/pytorch/ops/fused/grouped_mlp.py

@@ -1491,6 +1491,7 @@ def fuser_forward(
                     torch.int64,
                 ],
                 bulk_allocate=True,
+                split_alignment=256,
             )
 
         # Prepare FC1 grouped weight tensor for fused kernels.