[CUDA] Add opt-in direct scales for FP16 group-32 MatMulNBits M=1

.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

.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

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 "$<$<COMPILE_LANGUAGE:CXX>:/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
+        "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:-Xcompiler /permissive>"
+        "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:-Xcompiler /Zc:__cplusplus>"
+        "$<$<COMPILE_LANGUAGE:CUDA>: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
+        "$<$<COMPILE_LANGUAGE:CUDA>:--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 $<TARGET_FILE:${onnxruntime_provider_test_target}>
+          "--gtest_filter=${matmul4bits_m1_filter}"
+        WORKING_DIRECTORY $<TARGET_FILE_DIR:${onnxruntime_provider_test_target}>)
+      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})

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 <class T>
 bool TryMatMul4BitsM1(
     T* output,

onnxruntime/contrib_ops/cuda/quantization/matmul_4bits_m1_impl.cuh

@@ -5,16 +5,21 @@
 
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
+#include <cstdio>
+#include <type_traits>
 
 #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 <class T, int block_size, bool has_zero_point, bool TailN = false>
+template <class T, int block_size, bool has_zero_point, bool TailN = false,
+          bool direct_scales = false, int max_unroll = 16>
 __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<T, half> && 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<T*>(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<uint8_t*>(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<uint8_t*>(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<bool>("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 <class T>
 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<T, half>) {
+    // Cache the mode to keep launches consistent with captured graphs.
+    static const int experiment = ParseEnvironmentVariableWithDefault<int>("ORT_MATMUL4BITS_M1_EXPERIMENT", 0);
+    const int direct_unroll = SelectMatMul4BitsM1DirectUnroll(experiment, n, k, block_size, zero_points != nullptr);
+    if (direct_unroll == 4) {
+      MatMulFloat4BitsKernelM1<T, 32, false, false, true, 4><<<blocks, threads, 0, stream>>>(
+          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<T, 32, false, false, true, 16><<<blocks, threads, 0, stream>>>(
+          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<T, bs, true><<<blocks, threads, shared_mem_size, stream>>>(  \
@@ -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;
 }
 

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

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 <string>
+
+#include <cuda_runtime_api.h>
+
+#include "gtest/gtest.h"
+
+namespace onnxruntime::cuda::test {
+
+std::string CompareMatMul4BitsM1DirectScales(int max_unroll, int k, int n, int pattern);
+
+namespace {
+
+template <int MaxUnroll>
+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