[CUDA] MatMul: add an opt-in split-K GEMV for small-N fp16 shapes

onnxruntime/core/providers/cuda/math/matmul.cc

@@ -3,15 +3,21 @@
 
 #include "core/providers/cuda/math/matmul.h"
 
+#include "core/platform/env_var_utils.h"
 #include "core/providers/cuda/shared_inc/fpgeneric.h"
 #include "core/providers/cuda/cuda_allocator.h"
+#include "core/providers/cuda/math/matmul_small_n_gemv.h"
 #ifndef BUILD_CUDA_EP_AS_PLUGIN
 #include "core/providers/cuda/tunable/math/matmul.h"
 #endif
 
 namespace onnxruntime {
 namespace cuda {
 
+bool SmallNGemvEnabledFromEnvironment() {
+  return ParseEnvironmentVariableWithDefault<bool>("ORT_ENABLE_SMALL_N_GEMV", false);
+}
+
 #define REGISTER_KERNEL_TYPED(T)                                  \
   ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(                        \
       MatMul,                                                     \
@@ -322,6 +328,28 @@ Status MatMul<T>::ComputeDefault(OpKernelContext* ctx, MatMulComputeHelper& help
   auto& device_prop = GetDeviceProp();
 
   if (helper.OutputOffsets().size() == 1) {
+    if constexpr (std::is_same<T, MLFloat16>::value) {
+      // cuBLAS tiles this class of shape onto a handful of CTAs and spends
+      // several microseconds on a few hundred KiB of weights.
+      if (small_n_gemv_enabled_ && !transa && !transb && alpha_ == 1.0f &&
+          static_cast<int64_t>(lda) == helper.K() && static_cast<int64_t>(ldb) == helper.N() &&
+          static_cast<int64_t>(ldc) == helper.N() &&
+          CanUseSmallNGemv(helper.M(), helper.N(), helper.K(), left_X->DataRaw(), right_X->DataRaw(),
+                           Y->MutableDataRaw())) {
+        const int m = static_cast<int>(helper.M());
+        const int n = static_cast<int>(helper.N());
+        const int k = static_cast<int>(helper.K());
+        const size_t counter_elements = SmallNGemvCounterElements(n);
+        auto counter = GetScratchBuffer<unsigned int>(counter_elements, this->GetComputeStream(ctx));
+        CUDA_RETURN_IF_ERROR(cudaMemsetAsync(counter.get(), 0, counter_elements * sizeof(unsigned int), Stream(ctx)));
+        auto workspace = GetScratchBuffer<float>(SmallNGemvWorkspaceElements(m, n, k), this->GetComputeStream(ctx));
+        return LaunchSmallNGemv(Stream(ctx),
+                                reinterpret_cast<const half*>(left_X->Data<T>()),
+                                reinterpret_cast<const half*>(right_X->Data<T>()),
+                                reinterpret_cast<half*>(Y->MutableData<T>()),
+                                m, n, k, workspace.get(), counter.get());
+      }
+    }
     CUBLAS_RETURN_IF_ERROR(cublasGemmHelper(
         GetCublasHandle(ctx),
         transB,

onnxruntime/core/providers/cuda/math/matmul.h

@@ -8,6 +8,9 @@
 
 namespace onnxruntime {
 namespace cuda {
+
+bool SmallNGemvEnabledFromEnvironment();
+
 template <typename T>
 class MatMul final : public CudaKernel {
   using Base = CudaKernel;
@@ -19,7 +22,8 @@ class MatMul final : public CudaKernel {
         trans_A_{info.GetAttrOrDefault<int64_t>("transA", 0) != 0},
         trans_B_{info.GetAttrOrDefault<int64_t>("transB", 0) != 0},
         trans_batch_a_{info.GetAttrOrDefault<int64_t>("transBatchA", 0) != 0},
-        trans_batch_b_{info.GetAttrOrDefault<int64_t>("transBatchB", 0) != 0} {}
+        trans_batch_b_{info.GetAttrOrDefault<int64_t>("transBatchB", 0) != 0},
+        small_n_gemv_enabled_{SmallNGemvEnabledFromEnvironment()} {}
 
   Status ComputeInternal(OpKernelContext* context) const override;
   Status ComputeDefault(OpKernelContext* context, MatMulComputeHelper& helper) const;
@@ -30,6 +34,7 @@ class MatMul final : public CudaKernel {
   const bool trans_B_;
   const bool trans_batch_a_;
   const bool trans_batch_b_;
+  const bool small_n_gemv_enabled_;
 };
 
 template <typename T>

onnxruntime/core/providers/cuda/math/matmul_small_n_gemv.cu

@@ -0,0 +1,155 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "core/providers/cuda/math/matmul_small_n_gemv.h"
+
+#include "core/providers/cuda/cu_inc/common.cuh"
+
+namespace onnxruntime {
+namespace cuda {
+namespace {
+
+constexpr int kMaxM = 8;
+constexpr int kMaxN = 1024;
+constexpr int kMinK = 128;
+constexpr int kThreads = 256;
+constexpr int kTx = 32;  // one warp-wide column tile -> fully coalesced B reads
+constexpr int kMaxSplitK = 32;
+
+// grid = (ceil(n / kTx), split_k). Block (x, y) accumulates the K-slice
+// [y * k_per, (y+1) * k_per) for columns [x * kTx, x * kTx + kTx) and stores the
+// fp32 partials in `ws`. The last block to finish a column tile reduces the
+// partials in slice order (deterministic) and writes the half output.
+template <int M>
+__global__ void SmallNGemvSplitKKernel(const half* __restrict__ a, const half* __restrict__ b,
+                                       half* __restrict__ c, int n, int k,
+                                       volatile float* __restrict__ ws, unsigned int* __restrict__ counter,
+                                       int split_k) {
+  constexpr int TY = kThreads / kTx;
+  const int tx = static_cast<int>(threadIdx.x) % kTx;
+  const int ty = static_cast<int>(threadIdx.x) / kTx;
+  const int col = static_cast<int>(blockIdx.x) * kTx + tx;
+  const bool active = col < n;
+
+  const int k_per = (k + split_k - 1) / split_k;
+  const int k0 = static_cast<int>(blockIdx.y) * k_per;
+  const int k1 = min(k, k0 + k_per);
+
+  float acc[M];
+#pragma unroll
+  for (int m = 0; m < M; ++m) acc[m] = 0.0f;
+
+  if (active) {
+    for (int kk = k0 + ty; kk < k1; kk += TY) {
+      const float bv = __half2float(b[static_cast<size_t>(kk) * n + col]);
+#pragma unroll
+      for (int m = 0; m < M; ++m) {
+        acc[m] = fmaf(__half2float(a[static_cast<size_t>(m) * k + kk]), bv, acc[m]);
+      }
+    }
+  }
+
+  __shared__ float smem[kThreads * M];
+#pragma unroll
+  for (int m = 0; m < M; ++m) smem[threadIdx.x * M + m] = acc[m];
+  __syncthreads();
+#pragma unroll
+  for (int s = TY / 2; s > 0; s >>= 1) {
+    if (ty < s) {
+      const int partner = ((ty + s) * kTx + tx) * M;
+#pragma unroll
+      for (int m = 0; m < M; ++m) smem[threadIdx.x * M + m] += smem[partner + m];
+    }
+    __syncthreads();
+  }
+
+  if (ty == 0 && active) {
+#pragma unroll
+    for (int m = 0; m < M; ++m) {
+      ws[(static_cast<size_t>(blockIdx.y) * M + m) * n + col] = smem[threadIdx.x * M + m];
+    }
+  }
+
+  // Make the partials visible before announcing this block is done.
+  __threadfence();
+  __syncthreads();
+
+  __shared__ bool is_last;
+  if (threadIdx.x == 0) {
+    is_last = (atomicAdd(&counter[blockIdx.x], 1u) == static_cast<unsigned int>(split_k - 1));
+  }
+  __syncthreads();
+  if (!is_last) return;
+
+  if (ty == 0 && active) {
+#pragma unroll
+    for (int m = 0; m < M; ++m) {
+      float sum = 0.0f;
+      for (int s = 0; s < split_k; ++s) sum += ws[(static_cast<size_t>(s) * M + m) * n + col];
+      c[static_cast<size_t>(m) * n + col] = __float2half(sum);
+    }
+  }
+}
+
+template <int M>
+Status Launch(cudaStream_t stream, const half* a, const half* b, half* c, int n, int k,
+              float* ws, unsigned int* counter, int split_k) {
+  const dim3 grid(static_cast<unsigned>((n + kTx - 1) / kTx), static_cast<unsigned>(split_k));
+  SmallNGemvSplitKKernel<M><<<grid, kThreads, 0, stream>>>(a, b, c, n, k, ws, counter, split_k);
+  return CUDA_CALL(cudaGetLastError());
+}
+
+}  // namespace
+
+int SmallNGemvSplitK(int n, int k) {
+  const int tiles = (n + kTx - 1) / kTx;
+  int split_k = 1;
+  while (split_k < kMaxSplitK && split_k * tiles < 128) split_k <<= 1;
+  // Every slice must own at least one warp's worth of K rows.
+  while (split_k > 1 && k / split_k < kThreads / kTx) split_k >>= 1;
+  return split_k;
+}
+
+size_t SmallNGemvWorkspaceElements(int m, int n, int k) {
+  return static_cast<size_t>(SmallNGemvSplitK(n, k)) * m * n;
+}
+
+size_t SmallNGemvCounterElements(int n) {
+  return static_cast<size_t>((n + kTx - 1) / kTx);
+}
+
+bool CanUseSmallNGemv(int64_t m, int64_t n, int64_t k, const void* a, const void* b, const void* c) {
+  if (m < 1 || m > kMaxM || n < 1 || n > kMaxN || k < kMinK) return false;
+  if (k > (1 << 20)) return false;
+  const uintptr_t align = reinterpret_cast<uintptr_t>(a) | reinterpret_cast<uintptr_t>(b) |
+                          reinterpret_cast<uintptr_t>(c);
+  return (align % 8) == 0;
+}
+
+Status LaunchSmallNGemv(cudaStream_t stream, const half* a, const half* b, half* c,
+                        int m, int n, int k, float* ws, unsigned int* counter) {
+  const int split_k = SmallNGemvSplitK(n, k);
+  switch (m) {
+    case 1:
+      return Launch<1>(stream, a, b, c, n, k, ws, counter, split_k);
+    case 2:
+      return Launch<2>(stream, a, b, c, n, k, ws, counter, split_k);
+    case 3:
+      return Launch<3>(stream, a, b, c, n, k, ws, counter, split_k);
+    case 4:
+      return Launch<4>(stream, a, b, c, n, k, ws, counter, split_k);
+    case 5:
+      return Launch<5>(stream, a, b, c, n, k, ws, counter, split_k);
+    case 6:
+      return Launch<6>(stream, a, b, c, n, k, ws, counter, split_k);
+    case 7:
+      return Launch<7>(stream, a, b, c, n, k, ws, counter, split_k);
+    case 8:
+      return Launch<8>(stream, a, b, c, n, k, ws, counter, split_k);
+    default:
+      return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "SmallNGemv: unsupported M ", m);
+  }
+}
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/core/providers/cuda/math/matmul_small_n_gemv.h

@@ -0,0 +1,35 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "core/providers/cuda/cuda_common.h"
+
+namespace onnxruntime {
+namespace cuda {
+
+// Returns true when the shape/layout is eligible for LaunchSmallNGemv.
+bool CanUseSmallNGemv(int64_t m, int64_t n, int64_t k, const void* a, const void* b, const void* c);
+
+// Number of K-slices the launch will use, and the derived scratch sizes the
+// caller has to provide to LaunchSmallNGemv.
+int SmallNGemvSplitK(int n, int k);
+size_t SmallNGemvWorkspaceElements(int m, int n, int k);
+size_t SmallNGemvCounterElements(int n);
+
+// C[m, n] = sum_k A[m, k] * B[k, n], all row-major half, no transpose.
+// Targets decode-time projections whose N is far too small to keep a cuBLAS
+// tile kernel busy (router 2048x256, linear-attention gates 2048x32,
+// shared-expert gate 2048x1). The K axis is split across `SmallNGemvSplitK`
+// blocks so the read of B spreads over many SMs; the last block to finish a
+// column tile reduces the fp32 partials in slice order, so the result is
+// deterministic.
+//
+// `workspace` must hold SmallNGemvWorkspaceElements() floats. `counter` must
+// hold SmallNGemvCounterElements() unsigned ints and be zeroed before each
+// launch. Both buffers are exclusive to the launch.
+Status LaunchSmallNGemv(cudaStream_t stream, const half* a, const half* b, half* c,
+                        int m, int n, int k, float* workspace, unsigned int* counter);
+
+}  // namespace cuda
+}  // namespace onnxruntime

onnxruntime/test/providers/cuda/test_cases/matmul_small_n_gemv_op_test.cc

@@ -0,0 +1,58 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "gtest/gtest.h"
+
+#include <vector>
+
+#include "core/common/float16.h"
+#include "test/providers/provider_test_utils.h"
+#include "test/util/include/default_providers.h"
+#include "test/util/include/scoped_env_vars.h"
+
+namespace onnxruntime {
+namespace test {
+namespace {
+
+void RunMatMulOperatorCase(int m, int n, int k) {
+  std::vector<MLFloat16> a(static_cast<size_t>(m) * k);
+  std::vector<MLFloat16> b(static_cast<size_t>(k) * n);
+  std::vector<MLFloat16> expected(static_cast<size_t>(m) * n);
+  for (size_t index = 0; index < a.size(); ++index) {
+    a[index] = MLFloat16(static_cast<float>((index * 7 + 1) % 23) / 16.0f - 0.6875f);
+  }
+  for (size_t index = 0; index < b.size(); ++index) {
+    b[index] = MLFloat16(static_cast<float>((index * 11 + 2) % 19) / 16.0f - 0.5625f);
+  }
+  for (int row = 0; row < m; ++row) {
+    for (int col = 0; col < n; ++col) {
+      float sum = 0.0f;
+      for (int kk = 0; kk < k; ++kk) {
+        sum += a[static_cast<size_t>(row) * k + kk].ToFloat() *
+               b[static_cast<size_t>(kk) * n + col].ToFloat();
+      }
+      expected[static_cast<size_t>(row) * n + col] = MLFloat16(sum);
+    }
+  }
+
+  ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{{"ORT_ENABLE_SMALL_N_GEMV", "1"}}};
+  OpTester test("MatMul", 14);
+  test.AddInput<MLFloat16>("A", {m, k}, a);
+  test.AddInput<MLFloat16>("B", {k, n}, b);
+  test.AddOutput<MLFloat16>("Y", {m, n}, expected);
+  test.SetOutputAbsErr("Y", 0.05f);
+  test.ConfigEp(DefaultCudaExecutionProvider()).RunWithConfig();
+}
+
+TEST(MatMulSmallNGemvOpTest, DispatchesEligibleShapesWhenEnabled) {
+  RunMatMulOperatorCase(8, 1, 128);
+  RunMatMulOperatorCase(8, 1024, 128);
+}
+
+TEST(MatMulSmallNGemvOpTest, FallsBackForIneligibleShapeWhenEnabled) {
+  RunMatMulOperatorCase(8, 32, 127);
+}
+
+}  // namespace
+}  // namespace test
+}  // namespace onnxruntime