Add FP8 GEMV KSplit32 scheduling for the 48-SM SM121 GPU

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu

@@ -2,6 +2,7 @@
 // Licensed under the MIT License.
 
 #include "contrib_ops/cuda/math/matmul_block_scaled_fp8.h"
+#include "contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h"
 
 #include <cuda_bf16.h>
 #include <cuda_fp16.h>
@@ -852,19 +853,53 @@ int MatMulBlockScaledFp8GemvMaxM(int k, int block_size, const cudaDeviceProp& de
 #endif
 }
 
-Status LaunchMatMulBlockScaledFp8Gemv(void* y,
-                                      const void* a,
-                                      const void* b_fp8,
-                                      const float* weight_scale,
-                                      const void* bias,
-                                      const float* act_scale,
-                                      int m,
-                                      int n,
-                                      int k,
-                                      int block_size,
-                                      bool is_bf16,
-                                      const cudaDeviceProp& device_prop,
-                                      cudaStream_t stream) {
+int ApplyFp8MmaKSplitOverride(int k_split, int m, int n, int k) {
+  static int const override_k_split =
+      onnxruntime::ParseEnvironmentVariableWithDefault<int>("ORT_FP8_GEMV_KSPLIT", 0);
+  static int const match_n =
+      onnxruntime::ParseEnvironmentVariableWithDefault<int>("ORT_FP8_GEMV_MATCH_N", 0);
+  static int const match_k =
+      onnxruntime::ParseEnvironmentVariableWithDefault<int>("ORT_FP8_GEMV_MATCH_K", 0);
+  ORT_ENFORCE(override_k_split == 0 || override_k_split == 4 || override_k_split == 8 ||
+                  override_k_split == 16 || override_k_split == 32,
+              "ORT_FP8_GEMV_KSPLIT must be 0, 4, 8, 16, or 32.");
+  ORT_ENFORCE(match_n >= 0 && match_k >= 0,
+              "ORT_FP8_GEMV_MATCH_N and ORT_FP8_GEMV_MATCH_K must be non-negative.");
+
+  if ((match_n != 0 && n != match_n) || (match_k != 0 && k != match_k) ||
+      override_k_split == 0) {
+    return k_split;
+  }
+  ORT_ENFORCE(override_k_split != 32 || m <= 8,
+              "ORT_FP8_GEMV_KSPLIT=32 supports M up to 8, got M=", m, ".");
+  return override_k_split;
+}
+
+bool Fp8MmaGb10TuningEnabled() {
+  static bool const enabled = [] {
+    const int disable_tuning =
+        onnxruntime::ParseEnvironmentVariableWithDefault<int>("ORT_FP8_GEMV_DISABLE_GB10_TUNING", 0);
+    ORT_ENFORCE(disable_tuning == 0 || disable_tuning == 1,
+                "ORT_FP8_GEMV_DISABLE_GB10_TUNING must be 0 or 1.");
+    return disable_tuning == 0;
+  }();
+  return enabled;
+}
+
+static Status LaunchMatMulBlockScaledFp8GemvImpl(void* y,
+                                                 const void* a,
+                                                 const void* b_fp8,
+                                                 const float* weight_scale,
+                                                 const void* bias,
+                                                 const float* act_scale,
+                                                 int m,
+                                                 int n,
+                                                 int k,
+                                                 int block_size,
+                                                 bool is_bf16,
+                                                 const cudaDeviceProp& device_prop,
+                                                 cudaStream_t stream,
+                                                 bool enable_gb10_ksplit32) {
 #if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080
   if (m <= 0 || n <= 0 || k <= 0) {
     return Status::OK();
@@ -884,14 +919,14 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
                       "MatMulBlockQuantizedFp8Weight GEMV supports M above ", kFp8MmaGemvTileM,
                       " only on the mma sub-path, got M=", m, ".");
     const size_t element_size = is_bf16 ? sizeof(__nv_bfloat16) : sizeof(half);
-    ORT_RETURN_IF_ERROR(LaunchMatMulBlockScaledFp8Gemv(
+    ORT_RETURN_IF_ERROR(LaunchMatMulBlockScaledFp8GemvImpl(
         y, a, b_fp8, weight_scale, bias, act_scale, kFp8MmaGemvTileM, n, k, block_size,
-        is_bf16, device_prop, stream));
-    return LaunchMatMulBlockScaledFp8Gemv(
+        is_bf16, device_prop, stream, false));
+    return LaunchMatMulBlockScaledFp8GemvImpl(
         static_cast<uint8_t*>(y) + static_cast<size_t>(kFp8MmaGemvTileM) * n * element_size,
         static_cast<const uint8_t*>(a) + static_cast<size_t>(kFp8MmaGemvTileM) * k * element_size,
         b_fp8, weight_scale, bias, act_scale, m - kFp8MmaGemvTileM, n, k, block_size,
-        is_bf16, device_prop, stream);
+        is_bf16, device_prop, stream, false);
   }
 
   // Tensor-core path (SM80+). Beats the FMA kernel at every M on H200: 1.06-1.23x at M == 1 and
@@ -905,10 +940,13 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
   if (device_prop.major >= 8 && m <= kFp8MmaGemvTileM && k % 64 == 0 && k >= 256 &&
       block_size % 64 == 0 && Fp8GemvMmaEnabled()) {
     const int windows = k / 64;
-    int k_split = (n >= 8192) ? 8 : 16;  // wide N already fills the grid, so fewer warps per block
-    if (windows < k_split) {
-      k_split = (windows >= 8) ? 8 : 4;
-    }
+    // Preserve the generic schedule for recursive tiles from requests above the qualified M range.
+    const int selected_k_split =
+        enable_gb10_ksplit32 && Fp8MmaGb10TuningEnabled()
+            ? PickFp8MmaKSplit(n, m, windows, device_prop.multiProcessorCount,
+                               device_prop.major, device_prop.minor)
+            : PickGenericFp8MmaKSplit(n, windows);
+    const int k_split = ApplyFp8MmaKSplitOverride(selected_k_split, m, n, k);
     const int mtiles = (m > 16) ? 4 : ((m > 8) ? 2 : 1);
     const dim3 mma_blocks{static_cast<unsigned int>((n + 15) / 16)};
     const auto launch_mma = [&]<int KSplit, int MTiles>() {
@@ -934,7 +972,10 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
         launch_mma.template operator()<KSplit, 4>();
       }
     };
-    if (k_split == 16) {
+    if (k_split == 32) {
+      ORT_ENFORCE(mtiles == 1, "FP8 GEMV KSplit32 supports only M up to 8.");
+      launch_mma.template operator()<32, 1>();
+    } else if (k_split == 16) {
       launch_for_ksplit.template operator()<16>();
     } else if (k_split == 8) {
       launch_for_ksplit.template operator()<8>();
@@ -1023,8 +1064,27 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
   ORT_UNUSED_PARAMETER(is_bf16);
   ORT_UNUSED_PARAMETER(device_prop);
   ORT_UNUSED_PARAMETER(stream);
+  ORT_UNUSED_PARAMETER(enable_gb10_ksplit32);
   return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MatMulBlockQuantizedFp8Weight requires CUDA 11.8 or later.");
 #endif
 }
 
+Status LaunchMatMulBlockScaledFp8Gemv(void* y,
+                                      const void* a,
+                                      const void* b_fp8,
+                                      const float* weight_scale,
+                                      const void* bias,
+                                      const float* act_scale,
+                                      int m,
+                                      int n,
+                                      int k,
+                                      int block_size,
+                                      bool is_bf16,
+                                      const cudaDeviceProp& device_prop,
+                                      cudaStream_t stream) {
+  return LaunchMatMulBlockScaledFp8GemvImpl(
+      y, a, b_fp8, weight_scale, bias, act_scale, m, n, k, block_size,
+      is_bf16, device_prop, stream, true);
+}
+
 }  // namespace onnxruntime::contrib::cuda

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h

@@ -0,0 +1,43 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+namespace onnxruntime::contrib::cuda {
+
+inline int PickGenericFp8MmaKSplit(int n, int windows) {
+  int k_split = (n >= 8192) ? 8 : 16;
+  if (windows < k_split) {
+    k_split = (windows >= 8) ? 8 : 4;
+  }
+  return k_split;
+}
+
+inline int PickFp8MmaKSplit(int n, int m, int windows, int sm_count,
+                            int compute_capability_major, int compute_capability_minor) {
+  int k_split = PickGenericFp8MmaKSplit(n, windows);
+
+  constexpr int kOutputColumnsPerBlock = 16;
+  constexpr int kWideOutputMinBlocks = 1024;
+  constexpr int kLongReductionMinBlocks = 320;
+  constexpr int kWideOutputMinWindows = 80;
+  constexpr int kLongReductionMinWindows = 128;
+  const int output_blocks = (n + kOutputColumnsPerBlock - 1) / kOutputColumnsPerBlock;
+
+  // The qualified 48-SM SM121 GPU benefits from KSplit32 in two measured low-M regimes:
+  // wide outputs with substantial K and narrower outputs with very long reductions.
+  // The wide regime remains beneficial through the measured N=248320 lm-head shape, so it
+  // has no upper bound. Express these SM121 thresholds as output blocks so shapes with
+  // identical launch geometry use the same override; leave the generic selector unchanged
+  // to preserve behavior on other devices.
+  if (compute_capability_major == 12 && compute_capability_minor == 1 &&
+      sm_count == 48 && m <= 8 &&
+      ((output_blocks >= kWideOutputMinBlocks && windows >= kWideOutputMinWindows) ||
+       (output_blocks >= kLongReductionMinBlocks && windows >= kLongReductionMinWindows))) {
+    k_split = 32;
+  }
+
+  return k_split;
+}
+
+}  // namespace onnxruntime::contrib::cuda

onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc

@@ -1,18 +1,28 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include <cstdlib>
+
 #include "gtest/gtest.h"
 #include "test/common/cuda_op_test_utils.h"
 #include "test/common/tensor_op_test_utils.h"
 #include "test/providers/provider_test_utils.h"
 #include "test/unittest_util/conversion.h"
 #include "test/util/include/scoped_env_vars.h"
 
+#ifdef _WIN32
+#include <windows.h>
+#else
+#include <limits.h>
+#include <unistd.h>
+#endif
+
 #if defined(USE_CUDA)
 // CUDA_VERSION comes from cuda.h. Without this include the guard below silently
 // evaluates to false and every test in this file is compiled out.
 #include <cuda.h>
 
+#include "contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h"
 #include "core/providers/cuda/cuda_provider_options.h"
 #endif
 
@@ -30,6 +40,22 @@ namespace onnxruntime::test {
 // Dequantized weight value is fp8_e4m3(B[n, k]) * b_scale[n, k / block_size].
 
 namespace {
+std::string CurrentExecutablePath() {
+#ifdef _WIN32
+  std::string path(MAX_PATH, '\0');
+  const DWORD length = GetModuleFileNameA(nullptr, path.data(), static_cast<DWORD>(path.size()));
+  ORT_ENFORCE(length != 0 && length < path.size(), "GetModuleFileNameA failed.");
+  path.resize(length);
+  return path;
+#else
+  std::string path(PATH_MAX, '\0');
+  const ssize_t length = readlink("/proc/self/exe", path.data(), path.size());
+  ORT_ENFORCE(length > 0 && static_cast<size_t>(length) < path.size(), "readlink(/proc/self/exe) failed.");
+  path.resize(static_cast<size_t>(length));
+  return path;
+#endif
+}
+
 // Builds a [N, K] FP8 E4M3 weight where every element of row r equals row_value[r].
 std::vector<Float8E4M3FN> MakeConstRowWeight(const std::vector<float>& row_value, int64_t k) {
   std::vector<Float8E4M3FN> b(static_cast<size_t>(row_value.size()) * static_cast<size_t>(k));
@@ -42,6 +68,147 @@ std::vector<Float8E4M3FN> MakeConstRowWeight(const std::vector<float>& row_value
 }
 }  // namespace
 
+TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreKSplitSelection) {
+  struct Case {
+    int n;
+    int m;
+    int windows;
+    int sm_count;
+    int compute_capability_major;
+    int compute_capability_minor;
+    int expected;
+  };
+  const Case cases[] = {
+      {17408, 1, 80, 48, 12, 1, 32},
+      {16384, 8, 80, 48, 12, 1, 32},
+      {5120, 8, 128, 48, 12, 1, 32},
+      {16369, 1, 80, 48, 12, 1, 32},
+      {16368, 1, 80, 48, 12, 1, 8},
+      {16384, 1, 79, 48, 12, 1, 8},
+      {5105, 1, 128, 48, 12, 1, 32},
+      {5104, 1, 128, 48, 12, 1, 16},
+      {5120, 1, 127, 48, 12, 1, 16},
+      {32768, 1, 80, 48, 12, 1, 32},
+      {32769, 1, 80, 48, 12, 1, 32},
+      {65536, 1, 80, 48, 12, 1, 32},
+      {131072, 1, 80, 48, 12, 1, 32},
+      {248320, 1, 80, 48, 12, 1, 32},
+      {1024, 4, 80, 48, 12, 1, 16},
+      {7168, 8, 80, 48, 12, 1, 16},
+      {16384, 9, 80, 48, 12, 1, 8},
+      {5120, 9, 128, 48, 12, 1, 16},
+      {16384, 16, 80, 48, 12, 1, 8},
+      {16384, 1, 80, 47, 12, 1, 8},
+      {16384, 1, 80, 49, 12, 1, 8},
+      {16384, 1, 80, 48, 12, 0, 8},
+      {16384, 1, 80, 48, 9, 0, 8},
+  };
+
+  for (const Case& c : cases) {
+    SCOPED_TRACE("N = " + std::to_string(c.n) +
+                 ", M = " + std::to_string(c.m) +
+                 ", windows = " + std::to_string(c.windows) +
+                 ", SMs = " + std::to_string(c.sm_count) +
+                 ", CC = " + std::to_string(c.compute_capability_major) + "." +
+                 std::to_string(c.compute_capability_minor));
+    EXPECT_EQ(onnxruntime::contrib::cuda::PickFp8MmaKSplit(
+                  c.n, c.m, c.windows, c.sm_count,
+                  c.compute_capability_major, c.compute_capability_minor),
+              c.expected);
+  }
+}
+
+TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreForcedKSplit32) {
+  constexpr const char* kChildProcessVariable = "ORT_FP8_GEMV_KSPLIT_TEST_CHILD";
+  const bool is_child_process = !Env::Default().GetEnvironmentVar(kChildProcessVariable).empty();
+  if (!HasCudaEnvironment(800)) {
+    GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight.";
+  }
+
+  ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{
+      {"ORT_FP8_GEMV_MMA", "1"},
+      {"ORT_FP8_GEMV_MAX_M", "32"},
+      {"ORT_FP8_GEMV_KSPLIT", "32"},
+      {"ORT_FP8_GEMV_MATCH_N", "17"},
+      {"ORT_FP8_GEMV_MATCH_K", "2112"},
+      {kChildProcessVariable, "1"},
+  }};
+  if (!is_child_process) {
+    const std::string command =
+        "\"" + CurrentExecutablePath() +
+        "\" --gtest_filter=MatMulBlockQuantizedFp8WeightOpTest.GemvTensorCoreForcedKSplit32 --gtest_color=no";
+    ASSERT_EQ(std::system(command.c_str()), 0);
+    return;
+  }
+
+  constexpr int64_t m = 8;
+  constexpr int64_t n = 17;
+  constexpr int64_t k = 2112;  // 33 windows exercise a ragged KSplit32 reduction.
+  constexpr int64_t block_size = 64;
+  constexpr int64_t k_blocks = k / block_size;
+
+  static const float kWeightValues[] = {1.0f, 2.0f, -1.0f};
+  static const float kActValues[] = {1.0f, -1.0f, 0.5f, -0.5f};
+  std::vector<Float8E4M3FN> b(static_cast<size_t>(n * k));
+  std::vector<float> b_ref(static_cast<size_t>(n * k));
+  for (int64_t col = 0; col < n; ++col) {
+    for (int64_t i = 0; i < k; ++i) {
+      const float value = kWeightValues[(col + i) % 3];
+      b[static_cast<size_t>(col * k + i)] = Float8E4M3FN(value);
+      b_ref[static_cast<size_t>(col * k + i)] = value;
+    }
+  }
+  std::vector<float> b_scale(static_cast<size_t>(n * k_blocks));
+  for (int64_t col = 0; col < n; ++col) {
+    for (int64_t kb = 0; kb < k_blocks; ++kb) {
+      b_scale[static_cast<size_t>(col * k_blocks + kb)] =
+          static_cast<float>(1 + (col + kb) % 3) / 4.0f;
+    }
+  }
+  std::vector<float> a(static_cast<size_t>(m * k));
+  for (int64_t row = 0; row < m; ++row) {
+    for (int64_t i = 0; i < k; ++i) {
+      a[static_cast<size_t>(row * k + i)] = kActValues[(row + i) % 4];
+    }
+  }
+  std::vector<float> expected(static_cast<size_t>(m * n));
+  for (int64_t row = 0; row < m; ++row) {
+    for (int64_t col = 0; col < n; ++col) {
+      float acc = 0.0f;
+      for (int64_t i = 0; i < k; ++i) {
+        acc += a[static_cast<size_t>(row * k + i)] * b_ref[static_cast<size_t>(col * k + i)] *
+               b_scale[static_cast<size_t>(col * k_blocks + i / block_size)];
+      }
+      expected[static_cast<size_t>(row * n + col)] = acc;
+    }
+  }
+
+  {
+    OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain);
+    test.AddAttribute<int64_t>("block_size", block_size);
+    test.AddInput<MLFloat16>("A", {m, k}, FloatsToMLFloat16s(a));
+    test.AddInput<Float8E4M3FN>("B", {n, k}, b);
+    test.AddInput<float>("b_scale", {n, k_blocks}, b_scale);
+    test.AddOutput<MLFloat16>("Y", {m, n}, FloatsToMLFloat16s(expected));
+    test.SetOutputTolerance(0.005f);
+    std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+    execution_providers.push_back(DefaultCudaExecutionProvider());
+    test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+  }
+  {
+    OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain);
+    test.AddAttribute<int64_t>("block_size", block_size);
+    test.AddInput<BFloat16>("A", {m, k}, FloatsToBFloat16s(a));
+    test.AddInput<Float8E4M3FN>("B", {n, k}, b);
+    test.AddInput<float>("b_scale", {n, k_blocks}, b_scale);
+    test.AddOutput<BFloat16>("Y", {m, n}, FloatsToBFloat16s(expected));
+    test.SetOutputTolerance(0.05f);
+    std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+    execution_providers.push_back(DefaultCudaExecutionProvider());
+    test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+  }
+}
+
 // GEMM path (K not a multiple of 16 forces the cuBLAS dequant path), FP16 activations.
 // Weights are constant per row, so Y[m, n] = W_val[n] * sum_k A[m, k].
 TEST(MatMulBlockQuantizedFp8WeightOpTest, WeightOnlyGemmFp16) {