diff --git a/.github/workflows/linux_cuda_plugin_ci.yml b/.github/workflows/linux_cuda_plugin_ci.yml index 2c394e7868da3..c4bdba98949a1 100644 --- a/.github/workflows/linux_cuda_plugin_ci.yml +++ b/.github/workflows/linux_cuda_plugin_ci.yml @@ -123,6 +123,9 @@ jobs: cd /build/Release/Release ./onnxruntime_test_all ./onnxruntime_provider_test + for mode in 0 1 2; do + ORT_MATMUL4BITS_M1_EXPERIMENT=\$mode ORT_MATMUL4BITS_M1_LOG_DISPATCH=1 ./onnxruntime_provider_test --gtest_filter=MatMulNBits.Fp16_Int4_DirectScaleExperiment:MatMulNBits.Cuda_DirectScaleExperimentFallbacks + done " - name: Run CUDA Plugin EP Python Tests diff --git a/.github/workflows/windows_cuda_plugin.yml b/.github/workflows/windows_cuda_plugin.yml index 804e8c38ed4da..dfea1be95955e 100644 --- a/.github/workflows/windows_cuda_plugin.yml +++ b/.github/workflows/windows_cuda_plugin.yml @@ -209,6 +209,15 @@ jobs: exit $lastExitCode } + $env:ORT_MATMUL4BITS_M1_LOG_DISPATCH = "1" + foreach ($mode in 0, 1, 2) { + $env:ORT_MATMUL4BITS_M1_EXPERIMENT = $mode + .\onnxruntime_provider_test.exe --gtest_filter=MatMulNBits.Fp16_Int4_DirectScaleExperiment:MatMulNBits.Cuda_DirectScaleExperimentFallbacks + if ($lastExitCode -ne 0) { + exit $lastExitCode + } + } + - name: Run CUDA Plugin EP Python Tests working-directory: ${{ github.workspace }}\onnxruntime\test\python\transformers shell: pwsh diff --git a/cmake/onnxruntime_unittests.cmake b/cmake/onnxruntime_unittests.cmake index 149c1cd74f239..51df190f17d1f 100644 --- a/cmake/onnxruntime_unittests.cmake +++ b/cmake/onnxruntime_unittests.cmake @@ -1030,7 +1030,9 @@ set(all_tests set(onnxruntime_test_providers_runtime_dependencies) if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND NOT onnxruntime_BUILD_CUDA_EP_AS_PLUGIN) if (NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_REDUCED_OPS_BUILD) - set(onnxruntime_test_cuda_kernels_src_patterns "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/*.cc") + set(onnxruntime_test_cuda_kernels_src_patterns + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/*.cc" + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/*.cu") endif() file(GLOB onnxruntime_test_providers_cuda_ut_src CONFIGURE_DEPENDS @@ -1040,7 +1042,9 @@ if (onnxruntime_ENABLE_CUDA_EP_INTERNAL_TESTS AND NOT onnxruntime_BUILD_CUDA_EP_ if (onnxruntime_DISABLE_CONTRIB_OPS OR onnxruntime_CUDA_MINIMAL) list(REMOVE_ITEM onnxruntime_test_providers_cuda_ut_src - "${TEST_SRC_DIR}/providers/cuda/test_cases/generation_cuda_impl_test.cc") + "${TEST_SRC_DIR}/providers/cuda/test_cases/generation_cuda_impl_test.cc" + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cc" + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cu") endif() # cuda_plugin_test_shims.cc provides onnxruntime::GetEnvironmentVar for the @@ -1283,7 +1287,7 @@ function(onnxruntime_apply_test_target_workarounds target) if (MSVC) # TODO: The test code for OpenVINO, QNN, and WebGPU is getting flagged with a warning from ABSL for unreachable code. # Need to figure out how those particular targets/build variants are failing, but regular windows is not. - target_compile_options(${target} PRIVATE "/wd4702") + target_compile_options(${target} PRIVATE "$<$:/wd4702>") endif() # TODO fix shorten-64-to-32 warnings @@ -1633,6 +1637,11 @@ endif() # These also have some support for dynamically specified plugin EPs. if (NOT onnxruntime_MINIMAL_BUILD AND NOT onnxruntime_REDUCED_OPS_BUILD) set(onnxruntime_provider_test_target onnxruntime_provider_test) +if(onnxruntime_USE_CUDA AND NOT onnxruntime_DISABLE_CONTRIB_OPS AND onnxruntime_BUILD_CUDA_EP_AS_PLUGIN) + list(APPEND onnxruntime_provider_test_srcs + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cc" + "${TEST_SRC_DIR}/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cu") +endif() if (WIN32 AND TARGET onnxruntime_providers_cuda_ut) set(onnxruntime_provider_test_target onnxruntime_provider_test_executable) endif() @@ -1677,6 +1686,37 @@ block() DEPENDS ${onnxruntime_provider_test_deps} ) + if (onnxruntime_USE_CUDA AND NOT onnxruntime_DISABLE_CONTRIB_OPS AND onnxruntime_BUILD_CUDA_EP_AS_PLUGIN) + ort_configure_cuda_cccl(${onnxruntime_provider_test_target}) + if (MSVC) + # Abseil headers need the same nvcc host-compiler settings as the provider. + target_compile_options(${onnxruntime_provider_test_target} PRIVATE + "$<$:SHELL:-Xcompiler /permissive>" + "$<$:SHELL:-Xcompiler /Zc:__cplusplus>" + "$<$:SHELL:-Xcompiler /Zc:preprocessor>") + endif() + if (CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.8) + # Match the provider workaround for overflow diagnostics in CUDA math headers. + target_compile_options(${onnxruntime_provider_test_target} PRIVATE + "$<$:--diag-suppress=221>") + endif() + endif() + + if (onnxruntime_USE_CUDA AND NOT onnxruntime_DISABLE_CONTRIB_OPS) + set(matmul4bits_m1_filter + "MatMulNBits.Fp16_Int4_DirectScaleExperiment:MatMulNBits.Cuda_DirectScaleExperimentFallbacks") + foreach(mode 0 1 2) + add_test(NAME onnxruntime_matmul4bits_m1_experiment_${mode} + COMMAND $ + "--gtest_filter=${matmul4bits_m1_filter}" + WORKING_DIRECTORY $) + set_tests_properties(onnxruntime_matmul4bits_m1_experiment_${mode} PROPERTIES + ENVIRONMENT "ORT_MATMUL4BITS_M1_EXPERIMENT=${mode};ORT_MATMUL4BITS_M1_LOG_DISPATCH=1" + FAIL_REGULAR_EXPRESSION "0 tests from 0 test suites" + RUN_SERIAL TRUE) + endforeach() + endif() + # Match the non-plugin WebGPU test-source selection above. if (onnxruntime_USE_WEBGPU AND NOT onnxruntime_USE_EP_API_ADAPTERS) target_include_directories(${onnxruntime_provider_test_target} PRIVATE ${WGSL_GENERATED_ROOT}) diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1.cuh index d5bd50956f7f6..634a2a236164d 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1.cuh @@ -12,6 +12,20 @@ namespace onnxruntime { namespace contrib { namespace cuda { +// Returns 0 for the staged kernel, or the direct-scale unroll factor. +constexpr int SelectMatMul4BitsM1DirectUnroll(int mode, int n, int k, int block_size, bool has_zero_point) { + if ((mode != 1 && mode != 2) || block_size != 32 || has_zero_point || + n <= 0 || n % 8 != 0 || k <= 0 || k % 8 != 0) { + return 0; + } + // Keep the staged kernel for the measured long-K, narrow-N regression. + if (mode == 2 && k >= 4096 && n < 1024) { + return 0; + } + // Four-way unrolling regressed long reductions. + return mode == 2 && k < 4096 ? 4 : 16; +} + template bool TryMatMul4BitsM1( T* output, diff --git a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh index fbab46da81b50..5505837579cb7 100644 --- a/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh @@ -5,16 +5,21 @@ #include #include +#include +#include #include "core/providers/cuda/cu_inc/common.cuh" #include "contrib_ops/cuda/quantization/matmul_4bits_common.cuh" #include "contrib_ops/cuda/quantization/matmul_4bits_m1.cuh" +// Include after common.cuh defines SHARED_PROVIDER. +#include "core/platform/env_var_utils.h" namespace onnxruntime { namespace contrib { namespace cuda { -template +template __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloat4BitsKernelM1( T* output, const T* a_data, @@ -32,44 +37,50 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloat4Bi const int n_id = n_block_id * kColsPerThreadBlock + warp_id; constexpr int k_per_iter = kWarpSize * kElementsPerThreadPerIteration; - extern __shared__ char shared_buffer[]; - T* b_scale_vec = (T*)shared_buffer; - int offset = n_block_id * kColsPerThreadBlock * blocks_per_K; - for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * blocks_per_K; - i += kColsPerThreadBlock * kWarpSize) { - if constexpr (TailN) { - const int remaining_cols = n - n_block_id * kColsPerThreadBlock; - const int valid_cols = remaining_cols < kColsPerThreadBlock ? remaining_cols : kColsPerThreadBlock; - b_scale_vec[i] = i < valid_cols * blocks_per_K ? scales_data[offset + i] : T(0.f); - } else { - b_scale_vec[i] = scales_data[offset + i]; + static_assert(!direct_scales || (std::is_same_v && block_size == 32 && !has_zero_point && !TailN)); + static_assert(max_unroll == 4 || max_unroll == 16); + const T* b_scale_vec; + uint8_t* b_zp_vec = nullptr; + if constexpr (direct_scales) { + b_scale_vec = scales_data + n_id * blocks_per_K; + } else { + extern __shared__ char shared_buffer[]; + T* staged_scales = reinterpret_cast(shared_buffer); + const int offset = n_block_id * kColsPerThreadBlock * blocks_per_K; + for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * blocks_per_K; + i += kColsPerThreadBlock * kWarpSize) { + if constexpr (TailN) { + const int remaining_cols = n - n_block_id * kColsPerThreadBlock; + const int valid_cols = remaining_cols < kColsPerThreadBlock ? remaining_cols : kColsPerThreadBlock; + staged_scales[i] = i < valid_cols * blocks_per_K ? scales_data[offset + i] : T(0.f); + } else { + staged_scales[i] = scales_data[offset + i]; + } } - } - uint8_t* b_zp_vec; - (void)b_zp_vec; - if constexpr (has_zero_point) { - b_zp_vec = reinterpret_cast(b_scale_vec + kColsPerThreadBlock * blocks_per_K); - const int b_zp_k = (blocks_per_K + 1) / 2; - int zp_offset = n_block_id * kColsPerThreadBlock * b_zp_k; - for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * b_zp_k; - i += kColsPerThreadBlock * kWarpSize) { - b_zp_vec[2 * i] = (zero_points[zp_offset + i] & 0x0f); - b_zp_vec[2 * i + 1] = (zero_points[zp_offset + i] >> 4); + if constexpr (has_zero_point) { + b_zp_vec = reinterpret_cast(staged_scales + kColsPerThreadBlock * blocks_per_K); + const int b_zp_k = (blocks_per_K + 1) / 2; + const int zp_offset = n_block_id * kColsPerThreadBlock * b_zp_k; + for (int i = warp_id * kWarpSize + lane_id; i < kColsPerThreadBlock * b_zp_k; + i += kColsPerThreadBlock * kWarpSize) { + b_zp_vec[2 * i] = (zero_points[zp_offset + i] & 0x0f); + b_zp_vec[2 * i + 1] = (zero_points[zp_offset + i] >> 4); + } + b_zp_vec += warp_id * b_zp_k * 2; } - b_zp_vec += warp_id * b_zp_k * 2; - } - __syncthreads(); - if constexpr (TailN) { - static_assert(!has_zero_point, "TailN is scale-only"); - // Every thread participates in the scale-loading barrier before an invalid output warp exits. - if (n_id >= n) { - return; + __syncthreads(); + if constexpr (TailN) { + static_assert(!has_zero_point, "TailN is scale-only"); + // Every thread participates in the scale-loading barrier before an invalid output warp exits. + if (n_id >= n) { + return; + } } + b_scale_vec = staged_scales + warp_id * blocks_per_K; } a_data += m_id * k + (lane_id << 3); - b_scale_vec += warp_id * blocks_per_K; T sums[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; int k_id = 0; @@ -96,8 +107,10 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloat4Bi } \ } while (false) - UnRollReduction(16); - UnRollReduction(4); + UnRollReduction(max_unroll); + if constexpr (max_unroll == 16) { + UnRollReduction(4); + } UnRollReduction(1); #undef UnRollReduction @@ -121,6 +134,16 @@ __global__ void __launch_bounds__(kWarpSize* kColsPerThreadBlock) MatMulFloat4Bi } } +// Opt-in receipts let operator tests verify launches across the provider boundary. +inline void LogMatMul4BitsM1Dispatch(int n, int k, int block_size, bool has_zero_point, int unroll) { + static const bool enabled = + ParseEnvironmentVariableWithDefault("ORT_MATMUL4BITS_M1_LOG_DISPATCH", false); + if (enabled && cudaPeekAtLastError() == cudaSuccess) { + std::fprintf(stderr, "[MatMul4BitsM1] K=%d N=%d block_size=%d has_zp=%d unroll=%d\n", + k, n, block_size, has_zero_point ? 1 : 0, unroll); + } +} + template bool TryMatMul4BitsM1( T* output, @@ -147,6 +170,24 @@ bool TryMatMul4BitsM1( dim3 blocks((n + kColsPerThreadBlock - 1) / kColsPerThreadBlock, 1); dim3 threads(onnxruntime::cuda::GPU_WARP_SIZE_HOST, kColsPerThreadBlock); + if constexpr (std::is_same_v) { + // Cache the mode to keep launches consistent with captured graphs. + static const int experiment = ParseEnvironmentVariableWithDefault("ORT_MATMUL4BITS_M1_EXPERIMENT", 0); + const int direct_unroll = SelectMatMul4BitsM1DirectUnroll(experiment, n, k, block_size, zero_points != nullptr); + if (direct_unroll == 4) { + MatMulFloat4BitsKernelM1<<>>( + output, a_data, b_data_quant, scales_data, nullptr, 1, n, k, blocks_per_K); + LogMatMul4BitsM1Dispatch(n, k, block_size, zero_points != nullptr, 4); + return true; + } + if (direct_unroll == 16) { + MatMulFloat4BitsKernelM1<<>>( + output, a_data, b_data_quant, scales_data, nullptr, 1, n, k, blocks_per_K); + LogMatMul4BitsM1Dispatch(n, k, block_size, zero_points != nullptr, 16); + return true; + } + } + #define MATMUL_FLOAT4B_M1_DISPATCH(bs) \ if (zero_points != nullptr) { \ MatMulFloat4BitsKernelM1<<>>( \ @@ -168,6 +209,7 @@ bool TryMatMul4BitsM1( ORT_THROW("block size ", block_size, " is not supported"); } #undef MATMUL_FLOAT4B_M1_DISPATCH + LogMatMul4BitsM1Dispatch(n, k, block_size, zero_points != nullptr, 0); return true; } diff --git a/onnxruntime/python/tools/transformers/requirements.txt b/onnxruntime/python/tools/transformers/requirements.txt index dd8c91ff6df6e..ee1e1aaa25159 100644 --- a/onnxruntime/python/tools/transformers/requirements.txt +++ b/onnxruntime/python/tools/transformers/requirements.txt @@ -6,8 +6,8 @@ py3nvml packaging transformers >= 5.10.1 scipy -sentencepiece -pillow +sentencepiece==0.2.2 +pillow >= 13.0.0 # please follow https://pytorch.org/ to install PyTorch for your OS torch >= 2.6.0 diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cc b/onnxruntime/test/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cc new file mode 100644 index 0000000000000..a8785be36009c --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cc @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include + +#include + +#include "gtest/gtest.h" + +namespace onnxruntime::cuda::test { + +std::string CompareMatMul4BitsM1DirectScales(int max_unroll, int k, int n, int pattern); + +namespace { + +template +void TestDirectScales() { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + GTEST_SKIP() << "CUDA device unavailable"; + } + const auto compare = [](int k, int n, int pattern) { + SCOPED_TRACE(::testing::Message() << "K=" << k << ", N=" << n << ", pattern=" << pattern); + EXPECT_EQ("", CompareMatMul4BitsM1DirectScales(MaxUnroll, k, n, pattern)); + }; + for (int k : {8, 264, 1032, 3072, 4088, 4096, 4104, 8192}) { + compare(k, 24, 0); + } + compare(3072, 1024, 0); + compare(3072, 24, 1); +} + +TEST(CUDA_EP_Unittest, MatMul4BitsM1DirectScalesUnroll4) { + TestDirectScales<4>(); +} + +TEST(CUDA_EP_Unittest, MatMul4BitsM1DirectScalesUnroll16) { + TestDirectScales<16>(); +} + +} // namespace +} // namespace onnxruntime::cuda::test diff --git a/onnxruntime/test/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cu b/onnxruntime/test/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cu new file mode 100644 index 0000000000000..5b2cae593f537 --- /dev/null +++ b/onnxruntime/test/contrib_ops/cuda_kernels/matmul_4bits_m1_test.cu @@ -0,0 +1,99 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include + +#include +#include + +#include "contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh" + +namespace onnxruntime { +namespace cuda { +namespace test { +namespace { + +template +std::string CompareDirectScales(int k, int n, int pattern) { + constexpr int group_size = 32; + const int groups = (k + group_size - 1) / group_size; + std::mt19937 random(1234 + pattern); + std::uniform_int_distribution activation(-8, 8); + std::uniform_int_distribution scale(1, 8); + std::uniform_int_distribution packed(0, 255); + std::vector a(k), scales(n * groups); + std::vector weights(n * groups * group_size / 2); + for (int i = 0; i < k; ++i) { + const float value = pattern == 0 ? activation(random) / 16.0f + : (i % 2 == 0 ? 1.0f : -1.0f) * (i % 7 + 1) / 8.0f; + a[i] = __float2half_rn(value); + } + for (auto& value : scales) { + value = __float2half_rn(scale(random) / 32.0f); + } + for (auto& value : weights) { + value = static_cast(packed(random)); + } + + thrust::device_vector device_a(a), device_scales(scales); + thrust::device_vector device_weights(weights); + thrust::device_vector staged(n), direct(n); + auto* staged_ptr = thrust::raw_pointer_cast(staged.data()); + auto* direct_ptr = thrust::raw_pointer_cast(direct.data()); + if (const auto error = cudaMemset(staged_ptr, 0xff, n * sizeof(half)); error != cudaSuccess) { + return cudaGetErrorString(error); + } + if (const auto error = cudaMemset(direct_ptr, 0xff, n * sizeof(half)); error != cudaSuccess) { + return cudaGetErrorString(error); + } + + const dim3 grid(n / 8); + const dim3 block(32, 8); + const auto* a_ptr = thrust::raw_pointer_cast(device_a.data()); + const auto* weights_ptr = thrust::raw_pointer_cast(device_weights.data()); + const auto* scales_ptr = thrust::raw_pointer_cast(device_scales.data()); + contrib::cuda::MatMulFloat4BitsKernelM1 + <<>>( + staged_ptr, a_ptr, weights_ptr, scales_ptr, nullptr, 1, n, k, groups); + if (const auto error = cudaGetLastError(); error != cudaSuccess) { + return cudaGetErrorString(error); + } + // Explicit specialization makes coverage independent of the process-cached experiment mode. + contrib::cuda::MatMulFloat4BitsKernelM1 + <<>>( + direct_ptr, a_ptr, weights_ptr, scales_ptr, nullptr, 1, n, k, groups); + if (const auto error = cudaGetLastError(); error != cudaSuccess) { + return cudaGetErrorString(error); + } + if (const auto error = cudaDeviceSynchronize(); error != cudaSuccess) { + return cudaGetErrorString(error); + } + + const thrust::host_vector expected(staged), actual(direct); + for (int i = 0; i < n; ++i) { + if (!std::isfinite(__half2float(expected[i])) || !std::isfinite(__half2float(actual[i]))) { + return "Non-finite output at index " + std::to_string(i); + } + } + return std::memcmp(expected.data(), actual.data(), n * sizeof(half)) == 0 ? "" : "Bitwise output mismatch"; +} + +} // namespace + +std::string CompareMatMul4BitsM1DirectScales(int max_unroll, int k, int n, int pattern) { + if (max_unroll == 4) { + return CompareDirectScales<4>(k, n, pattern); + } + if (max_unroll == 16) { + return CompareDirectScales<16>(k, n, pattern); + } + return "Unsupported unroll factor"; +} + +} // namespace test +} // namespace cuda +} // namespace onnxruntime diff --git a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc index 4d7434ef398dd..afb2290d10e7f 100644 --- a/onnxruntime/test/contrib_ops/matmul_4bits_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_4bits_test.cc @@ -5,6 +5,7 @@ #include #include +#include #include "gtest/gtest.h" #include "gmock/gmock.h" @@ -33,7 +34,11 @@ #include "core/providers/webgpu/webgpu_provider_options.h" #include "core/framework/prepacked_weights_container.h" #include "core/session/onnxruntime_session_options_config_keys.h" +#include "core/platform/env_var_utils.h" #include "test/util/include/test/test_environment.h" +#ifdef USE_CUDA +#include "contrib_ops/cuda/quantization/matmul_4bits_m1.cuh" +#endif extern std::unique_ptr ort_env; @@ -1678,6 +1683,119 @@ TEST(MatMulNBits, Fp16_Int4_PrepackedSm90BlockSize32Rejected) { // MatMulNBitsSm90ValidationTest in test/contrib_ops/cuda_kernels/matmul_nbits_sm90_validation_test.cc // (see the comment there for why that coverage lives in a separate translation unit). +#ifdef USE_CUDA +TEST(MatMulNBits, DirectScaleExperimentSelection) { + using contrib::cuda::SelectMatMul4BitsM1DirectUnroll; + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(0, 1024, 3072, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(1, 1024, 3072, 32, false), 16); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1024, 3072, 32, false), 4); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 8, 8, 32, false), 4); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 8, 4088, 32, false), 4); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 8, 4096, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 8, 8192, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 512, 4096, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1024, 4096, 32, false), 16); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1024, 8192, 32, false), 16); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(3, 1024, 3072, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1024, 3072, 64, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1024, 3072, 32, true), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1023, 3072, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1024, 3071, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 0, 3072, 32, false), 0); + EXPECT_EQ(SelectMatMul4BitsM1DirectUnroll(2, 1024, 0, 32, false), 0); +} + +// Use fresh processes for modes 0/1/2; the launcher caches the mode. +TEST(MatMulNBits, Fp16_Int4_DirectScaleExperiment) { + if (!DefaultCudaExecutionProvider()) { + GTEST_SKIP() << "CUDA execution provider unavailable"; + } + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_FPA_INTB_GEMM", "0"}}}; + const bool verify_dispatch = + ParseEnvironmentVariableWithDefault("ORT_MATMUL4BITS_M1_LOG_DISPATCH", false); + const int mode = ParseEnvironmentVariableWithDefault("ORT_MATMUL4BITS_M1_EXPERIMENT", 0); + for (auto k : {8, 24, 264, 1032, 3072, 4088, 4096, 8192}) { + for (auto n : {8, 24, 256, 1024}) { + TestOptions opts{}; + opts.M = 1, opts.N = n, opts.K = k; + opts.block_size = 32; + opts.disable_cpu_ep_fallback = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + if (verify_dispatch) { + testing::internal::CaptureStderr(); + } + RunTest(opts, std::move(eps)); + if (verify_dispatch) { + const int unroll = contrib::cuda::SelectMatMul4BitsM1DirectUnroll(mode, n, k, 32, false); + EXPECT_THAT(testing::internal::GetCapturedStderr(), + testing::HasSubstr(MakeString("[MatMul4BitsM1] K=", k, " N=", n, + " block_size=32 has_zp=0 unroll=", unroll, "\n"))); + } + } + } +} +#endif + +#ifdef USE_CUDA +TEST(MatMulNBits, Cuda_DirectScaleExperimentFallbacks) { + if (!DefaultCudaExecutionProvider()) { + GTEST_SKIP() << "CUDA execution provider unavailable"; + } + ScopedEnvironmentVariables env{EnvVarMap{{"ORT_FPA_INTB_GEMM", "0"}}}; + for (auto variant : {0, 1, 2, 3, 4, 5, 6}) { + TestOptions opts{}; + opts.M = 1, opts.N = 24, opts.K = 1032; + opts.block_size = 32; + opts.disable_cpu_ep_fallback = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + switch (variant) { + case 0: + opts.N = 23; + break; + case 1: + opts.K = 1031; + break; + case 2: + opts.M = 3; + break; + case 3: + opts.block_size = 16; + break; + case 4: + opts.block_size = 64; + break; + case 5: + opts.block_size = 128; + break; + case 6: + opts.has_zero_point = true; + break; + } + std::vector> eps; + eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(eps)); + } + TestOptions opts{}; + opts.M = 1, opts.N = 24, opts.K = 1032; + opts.block_size = 32; + opts.disable_cpu_ep_fallback = true; + opts.output_abs_error = 0.1f; + opts.output_rel_error = 0.02f; + std::vector> float_eps; + float_eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(float_eps)); + if (HasCudaEnvironment(800)) { + std::vector> bf16_eps; + bf16_eps.push_back(DefaultCudaExecutionProvider()); + RunTest(opts, std::move(bf16_eps)); + } +} +#endif + // Exercises the CUDA small-M batched GEMV tiles: CtaM in {2,4,8,16} (with M values that are not a // multiple of CtaM so the row-skip path runs) and CtaN in {1,2} (N divisible / not divisible by 16). TEST(MatMulNBits, Fp16_Int4_SmallMBatchedTiles) { diff --git a/tools/ci_build/requirements/transformers-test/requirements.txt b/tools/ci_build/requirements/transformers-test/requirements.txt index 9d096c235e592..0cb7209796ce9 100644 --- a/tools/ci_build/requirements/transformers-test/requirements.txt +++ b/tools/ci_build/requirements/transformers-test/requirements.txt @@ -8,7 +8,8 @@ torch==2.10.0 torchvision==0.25.0 transformers==5.10.1 parameterized>=0.8.1 -sentencepiece +# SentencePiece 0.2.3 requires protobuf 7.x, but CI uses protobuf 6.x. +sentencepiece==0.2.2 psutil einops onnxscript==0.6.2