From d258d7c7e1b06b8f37fa28fd6c2848c0e38c404a Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 4 Aug 2026 15:09:10 -0700 Subject: [PATCH 1/2] Harden CUDA fp16 transpose index math against int overflow Use int64 linear offsets in transposeNoOverlap, guard helper eligibility on m*n int32 fit, and make dim-grid math use int64 intermediates. Add CUDA unit tests for overflow and grid-y helper gating. --- onnxruntime/core/providers/cuda/fpgeneric.cu | 23 ++++++++++++++++--- .../cuda/test_cases/cuda_utils_test.cc | 11 +++++++++ 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/cuda/fpgeneric.cu b/onnxruntime/core/providers/cuda/fpgeneric.cu index 78426c1294883..7a82bff0c42f0 100644 --- a/onnxruntime/core/providers/cuda/fpgeneric.cu +++ b/onnxruntime/core/providers/cuda/fpgeneric.cu @@ -13,6 +13,8 @@ #include "core/providers/cuda/curand_wrapper.h" #include "core/providers/cuda/cu_inc/common.cuh" +#include + #define TRANS_TILE_DIM 32 #define BLOCK_ROWS 8 #define COPY_TILE_DIM 1024 @@ -31,7 +33,8 @@ __global__ void transposeNoOverlap(half* odata, const half* idata, const int m, if (x < m) { for (int j = 0; j < TRANS_TILE_DIM; j += BLOCK_ROWS) { if (j >= (n - y)) continue; - tile[threadIdx.y + j][threadIdx.x] = idata[(y + j) * m + x]; + const int64_t input_offset = static_cast(y + j) * m + x; + tile[threadIdx.y + j][threadIdx.x] = idata[input_offset]; } } @@ -44,7 +47,8 @@ __global__ void transposeNoOverlap(half* odata, const half* idata, const int m, for (int j = 0; j < TRANS_TILE_DIM; j += BLOCK_ROWS) { if ((y + j) >= m) return; - odata[(y + j) * n + x] = tile[threadIdx.x][threadIdx.y + j]; + const int64_t output_offset = static_cast(y + j) * n + x; + odata[output_offset] = tile[threadIdx.x][threadIdx.y + j]; } } @@ -64,11 +68,23 @@ __global__ void CopyVectorBFloat16(const onnxruntime::BFloat16* x, int incx, onn } // namespace dim3 cublasTransposeHelperDimGrid(int m, int n) { - return dim3((n + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, (m + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM, 1); + const auto grid_x = static_cast((static_cast(n) + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM); + const auto grid_y = static_cast((static_cast(m) + TRANS_TILE_DIM - 1) / TRANS_TILE_DIM); + return dim3(grid_x, grid_y, 1); } // cublasTransposeHelper can only be used if it won't overflow the 65536 grid y dimension size __host__ bool CanUse_cublasTransposeHelper_MLFloat16(int m, int n) { + if (m < 0 || n < 0) { + return false; + } + + // transposeNoOverlap uses row * stride + col addressing in device code. + // Keep fallback disabled when total element count would overflow int32 indexing. + if (static_cast(m) * static_cast(n) > std::numeric_limits::max()) { + return false; + } + dim3 dimGrid = cublasTransposeHelperDimGrid(m, n); return dimGrid.y < 65536; } @@ -78,6 +94,7 @@ cublasStatus_t cublasTransposeHelper(cudaStream_t stream, cublasHandle_t, cublas dim3 dimGrid = cublasTransposeHelperDimGrid(m, n); dim3 dimBlock(TRANS_TILE_DIM, BLOCK_ROWS, 1); + ORT_ENFORCE(static_cast(m) * static_cast(n) <= std::numeric_limits::max()); ORT_ENFORCE(dimGrid.y < 65536); // To prevent this, call CanUse_cublasTransposeHelper_MLFloat16 first transposeNoOverlap<<>>(C, A, n, m); } else { diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc index 7468a5718425e..ed12da155162f 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc @@ -8,6 +8,7 @@ #include "core/common/common.h" #include "core/providers/cuda/shared_inc/cuda_call.h" +#include "core/providers/cuda/shared_inc/fpgeneric.h" #include "core/providers/cuda/shared_inc/cuda_utils.h" namespace onnxruntime { @@ -48,6 +49,16 @@ TEST(CudaUtilsTest, FillCorrectness) { TestFillCorrectness(1 << 20, 5.0); } +TEST(CudaUtilsTest, CanUseTransposeHelperRejectsOverflowingElementCount) { + EXPECT_TRUE(CanUse_cublasTransposeHelper_MLFloat16(100, 100)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, 25000000)); +} + +TEST(CudaUtilsTest, CanUseTransposeHelperRejectsGridYOverflow) { + // For TRANS_TILE_DIM=32, m=2097152 yields grid_y=65536, which is out of range. + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(2097152, 1)); +} + } // namespace test } // namespace cuda } // namespace onnxruntime From 2f5e67725a0787a4f8d7afb4d7f078b757b31d18 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 5 Aug 2026 10:39:38 -0700 Subject: [PATCH 2/2] Address CUDA transpose overflow review comments Tighten the MLFloat16 transpose helper so it rejects non-positive dimensions, matching the valid launch configuration requirements and preventing zero-sized CUDA grids from being treated as usable. Also add an explicit precondition in the helper itself so invalid dimensions fail safely even if callers bypass the CanUse guard. Update the overflow comment to reflect the current int64_t indexing rationale and extend the CUDA utility tests to cover zero and negative dimensions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- onnxruntime/core/providers/cuda/fpgeneric.cu | 7 ++++--- .../test/providers/cuda/test_cases/cuda_utils_test.cc | 7 +++++++ 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/providers/cuda/fpgeneric.cu b/onnxruntime/core/providers/cuda/fpgeneric.cu index 7a82bff0c42f0..21097d1080f34 100644 --- a/onnxruntime/core/providers/cuda/fpgeneric.cu +++ b/onnxruntime/core/providers/cuda/fpgeneric.cu @@ -75,12 +75,12 @@ dim3 cublasTransposeHelperDimGrid(int m, int n) { // cublasTransposeHelper can only be used if it won't overflow the 65536 grid y dimension size __host__ bool CanUse_cublasTransposeHelper_MLFloat16(int m, int n) { - if (m < 0 || n < 0) { + if (m <= 0 || n <= 0) { return false; } - // transposeNoOverlap uses row * stride + col addressing in device code. - // Keep fallback disabled when total element count would overflow int32 indexing. + // transposeNoOverlap uses int64_t row * stride + col addressing in device code. + // Keep fallback disabled when total element count would overflow 32-bit launch/indexing assumptions. if (static_cast(m) * static_cast(n) > std::numeric_limits::max()) { return false; } @@ -90,6 +90,7 @@ __host__ bool CanUse_cublasTransposeHelper_MLFloat16(int m, int n) { } cublasStatus_t cublasTransposeHelper(cudaStream_t stream, cublasHandle_t, cublasOperation_t, cublasOperation_t, int m, int n, const half*, const half* A, int, const half*, const half*, int, half* C, int) { + ORT_ENFORCE(m > 0 && n > 0); if (C != A) { dim3 dimGrid = cublasTransposeHelperDimGrid(m, n); dim3 dimBlock(TRANS_TILE_DIM, BLOCK_ROWS, 1); diff --git a/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc b/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc index ed12da155162f..5cf06d02ca804 100644 --- a/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc +++ b/onnxruntime/test/providers/cuda/test_cases/cuda_utils_test.cc @@ -54,6 +54,13 @@ TEST(CudaUtilsTest, CanUseTransposeHelperRejectsOverflowingElementCount) { EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, 25000000)); } +TEST(CudaUtilsTest, CanUseTransposeHelperRejectsNonPositiveDimensions) { + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(0, 100)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, 0)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(-1, 100)); + EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(100, -1)); +} + TEST(CudaUtilsTest, CanUseTransposeHelperRejectsGridYOverflow) { // For TRANS_TILE_DIM=32, m=2097152 yields grid_y=65536, which is out of range. EXPECT_FALSE(CanUse_cublasTransposeHelper_MLFloat16(2097152, 1));