diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 4c0e844c8c4..6c06f393c66 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/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 diff --git a/transformer_engine/common/util/splits_to_offsets.cu b/transformer_engine/common/util/splits_to_offsets.cu index 2a06603dafc..169052fa9de 100644 --- a/transformer_engine/common/util/splits_to_offsets.cu +++ b/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 +__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(split) & alignment_mask) != 0; + } else { + misaligned = static_cast(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<<<1, s2o::kSmallThreadsPerBlock, 0, stream>>>(args); + } else { + s2o::kernel<<<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,6 +311,7 @@ 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; @@ -286,7 +319,11 @@ void nvte_splits_to_offsets_multi(NVTETensor split_sizes, NVTETensor *outputs, 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<<<1, s2o::kSmallThreadsPerBlock, 0, stream>>>(args); + } else { + s2o::kernel<<<1, s2o::kThreadsPerBlock, 0, stream>>>(args); + } NVTE_CHECK_CUDA(cudaGetLastError()); } } diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index f6cf7d2fa77..ef2bbc9b44a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/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> splits_to_offsets_multi( const at::Tensor &split_sizes, const c10::Device &device, const std::vector &strides, const std::vector &include_leading_zero, const std::vector &dtypes, - bool bulk_allocate_outputs); + bool bulk_allocate_outputs, size_t split_alignment); at::Tensor copy_data_ptrs_to_device(const std::vector &tensors, const c10::Device &device); diff --git a/transformer_engine/pytorch/csrc/extensions/misc.cpp b/transformer_engine/pytorch/csrc/extensions/misc.cpp index 01bcd7eb05a..e17e362438b 100644 --- a/transformer_engine/pytorch/csrc/extensions/misc.cpp +++ b/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> splits_to_offsets_multi( const at::Tensor &split_sizes, const c10::Device &device, const std::vector &strides, const std::vector &include_leading_zero, const std::vector &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(split_sizes.numel()); @@ -51,6 +51,7 @@ std::tuple> 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> 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)}; diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 89e7d92aca9..380bb6b9820 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/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()); diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 171500d7798..59b996743bc 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/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] diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py index d678d92749c..b0a90cffe5a 100644 --- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py +++ b/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.