Pin FP8 GEMV residency for grids just past two blocks per SM

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>
@@ -514,17 +515,17 @@ struct Fp8GemvMma<__nv_bfloat16> {
 };
 
 template <int KSplit, int MTiles, typename AType>
-__global__ void MatMulBlockScaledFp8MmaGemvKernel(AType* __restrict__ output,
-                                                  const AType* __restrict__ input_a,
-                                                  const __nv_fp8_e4m3* __restrict__ input_b,
-                                                  const float* __restrict__ weight_scale,
-                                                  const AType* __restrict__ bias,
-                                                  const float* __restrict__ act_scale,
-                                                  int m,
-                                                  int n,
-                                                  int k,
-                                                  int block_size,
-                                                  int k_blocks) {
+__device__ __forceinline__ void Fp8MmaGemvBody(AType* __restrict__ output,
+                                               const AType* __restrict__ input_a,
+                                               const __nv_fp8_e4m3* __restrict__ input_b,
+                                               const float* __restrict__ weight_scale,
+                                               const AType* __restrict__ bias,
+                                               const float* __restrict__ act_scale,
+                                               int m,
+                                               int n,
+                                               int k,
+                                               int block_size,
+                                               int k_blocks) {
   using Mma = Fp8GemvMma<AType>;
 
   const bool act_qdq = act_scale != nullptr;
@@ -691,12 +692,49 @@ __global__ void MatMulBlockScaledFp8MmaGemvKernel(AType* __restrict__ output,
   }
 }
 
+// Two entry points over one body. The pinned one carries a residency hint; see
+// `Fp8MmaGemvPinsResidency` for when the launcher picks it and why the plain one has to stay.
+// clang-format off
+#define ORT_FP8_MMA_GEMV_PARAMS                  \
+  AType* __restrict__ output,                    \
+      const AType* __restrict__ input_a,         \
+      const __nv_fp8_e4m3* __restrict__ input_b, \
+      const float* __restrict__ weight_scale,    \
+      const AType* __restrict__ bias,            \
+      const float* __restrict__ act_scale,       \
+      int m, int n, int k, int block_size, int k_blocks
+
+#define ORT_FP8_MMA_GEMV_ARGS \
+  output, input_a, input_b, weight_scale, bias, act_scale, m, n, k, block_size, k_blocks
+// clang-format on
+
+template <int KSplit, int MTiles, typename AType>
+__global__ void MatMulBlockScaledFp8MmaGemvKernel(ORT_FP8_MMA_GEMV_PARAMS) {
+  Fp8MmaGemvBody<KSplit, MTiles, AType>(ORT_FP8_MMA_GEMV_ARGS);
+}
+
+template <int KSplit, int MTiles, typename AType>
+__global__ __launch_bounds__(32 * KSplit, 3) void MatMulBlockScaledFp8MmaGemvKernelPinned(ORT_FP8_MMA_GEMV_PARAMS) {
+  Fp8MmaGemvBody<KSplit, MTiles, AType>(ORT_FP8_MMA_GEMV_ARGS);
+}
+
+#undef ORT_FP8_MMA_GEMV_ARGS
+#undef ORT_FP8_MMA_GEMV_PARAMS
+
 // Kill switch for A/B testing the tensor-core path against the FMA path in the same binary.
 bool Fp8GemvMmaEnabled() {
   static bool const enabled = onnxruntime::ParseEnvironmentVariableWithDefault<bool>("ORT_FP8_GEMV_MMA", true);
   return enabled;
 }
 
+// Tiling override for A/B sweeps; 0 keeps the heuristic.
+int Fp8GemvKSplitOverride() {
+  static int const k_split = onnxruntime::ParseEnvironmentVariableWithDefault<int>("ORT_FP8_GEMV_KSPLIT", 0);
+  ORT_ENFORCE(k_split == 0 || k_split == 4 || k_split == 8 || k_split == 16 || k_split == 32,
+              "ORT_FP8_GEMV_KSPLIT must be 0, 4, 8, 16, or 32.");
+  return k_split;
+}
+
 // Largest M each sub-path accepts. One mma launch unrolls 4 tiles of the mma's 8-row N extent.
 // Larger speculative batches are split into two launches so they keep the same per-row arithmetic
 // instead of switching to the dequantize + cuBLAS path.
@@ -905,23 +943,37 @@ 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;
+    int k_split = PickFp8MmaKSplit(n, m, windows, device_prop.multiProcessorCount, device_prop.major);
+    if (const int k_split_override = Fp8GemvKSplitOverride(); k_split_override != 0) {
+      k_split = k_split_override;
     }
     const int mtiles = (m > 16) ? 4 : ((m > 8) ? 2 : 1);
     const dim3 mma_blocks{static_cast<unsigned int>((n + 15) / 16)};
+    const bool pin_residency = Fp8MmaGemvPinsResidency(n, k_split, mtiles, device_prop.multiProcessorCount);
     const auto launch_mma = [&]<int KSplit, int MTiles>() {
       const dim3 mma_threads{32, KSplit};
-      if (is_bf16) {
-        MatMulBlockScaledFp8MmaGemvKernel<KSplit, MTiles><<<mma_blocks, mma_threads, 0, stream>>>(
-            reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b,
-            weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k, block_size, k_blocks);
-      } else {
-        MatMulBlockScaledFp8MmaGemvKernel<KSplit, MTiles><<<mma_blocks, mma_threads, 0, stream>>>(
-            reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,
-            weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k, block_size, k_blocks);
+#define ORT_FP8_LAUNCH_MMA(kernel_name)                                                      \
+  do {                                                                                       \
+    if (is_bf16) {                                                                           \
+      kernel_name<KSplit, MTiles><<<mma_blocks, mma_threads, 0, stream>>>(                   \
+          reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast<const __nv_bfloat16*>(a), b, \
+          weight_scale, reinterpret_cast<const __nv_bfloat16*>(bias), act_scale, m, n, k,    \
+          block_size, k_blocks);                                                             \
+    } else {                                                                                 \
+      kernel_name<KSplit, MTiles><<<mma_blocks, mma_threads, 0, stream>>>(                   \
+          reinterpret_cast<half*>(y), reinterpret_cast<const half*>(a), b,                   \
+          weight_scale, reinterpret_cast<const half*>(bias), act_scale, m, n, k,             \
+          block_size, k_blocks);                                                             \
+    }                                                                                        \
+  } while (0)
+      if constexpr (KSplit == 16 && MTiles == 1) {
+        if (pin_residency) {
+          ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernelPinned);
+          return;
+        }
       }
+      ORT_FP8_LAUNCH_MMA(MatMulBlockScaledFp8MmaGemvKernel);
+#undef ORT_FP8_LAUNCH_MMA
     };
     // Only 1, 2 and 4 row tiles are instantiated; an M of 17..24 rounds up to 4 and masks the
     // remainder, which costs nothing next to the weight traffic it shares.
@@ -934,7 +986,9 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y,
         launch_mma.template operator()<KSplit, 4>();
       }
     };
-    if (k_split == 16) {
+    if (k_split == 32) {
+      launch_for_ksplit.template operator()<32>();
+    } else if (k_split == 16) {
       launch_for_ksplit.template operator()<16>();
     } else if (k_split == 8) {
       launch_for_ksplit.template operator()<8>();

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h

@@ -0,0 +1,47 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+namespace onnxruntime::contrib::cuda {
+
+inline int PickFp8MmaKSplit(int n, int m, int windows, int sm_count, int compute_capability_major) {
+  int k_split = (n >= 8192) ? 8 : 16;
+  if (windows < k_split) {
+    k_split = (windows >= 8) ? 8 : 4;
+  }
+
+  // Low-SM-count client Blackwell GPUs benefit from additional K parallelism for low-M decode.
+  if (compute_capability_major == 12 && sm_count <= 64 && windows >= 16) {
+    k_split = 16;
+    if (m <= 8 && windows >= 32 && (n >= 4096 || m <= 2)) {
+      k_split = 32;
+    }
+  }
+
+  return k_split;
+}
+
+// True when the tensor-core GEMV should launch the entry point that carries a residency hint.
+//
+// The mma grid is ceil(N / 16) blocks. A 16-warp block only fits twice per SM, so N just above
+// 32 * sm_count spills into a second, nearly empty wave: on H200 N = 5120 launches 1.21 waves
+// and ncu measures 66% active cycles. __launch_bounds__(threads, 3) makes those shapes a single
+// wave, worth 1.21-1.35x. Outside that window it only costs registers, so:
+//
+//   * a grid at or below 2 blocks per SM is already one wave and must stay on the plain kernel;
+//   * a grid above 3 blocks per SM stays multi-wave either way;
+//   * 8-warp blocks (KSplit 8, taken from N >= 8192) must not carry the attribute at all --
+//     declaring it replaces nvcc's implicit bounds and costs 1.05-1.08x even when the register
+//     cap is unchanged, and KSplit 32 cannot host 3 blocks per SM at all;
+//   * only one row tile fits the 40-register cap that 3 blocks per SM imply. M = 16 (two tiles)
+//     measures 0.74x and M = 32 (four tiles) 0.24x, both from spills.
+inline bool Fp8MmaGemvPinsResidency(int n, int k_split, int m_tiles, int sm_count) {
+  if (k_split != 16 || m_tiles != 1) {
+    return false;
+  }
+  const int col_blocks = (n + 15) / 16;
+  return col_blocks > 2 * sm_count && col_blocks <= 3 * sm_count;
+}
+
+}  // namespace onnxruntime::contrib::cuda

onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc

@@ -12,7 +12,9 @@
 // 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 <cuda_runtime_api.h>
 
+#include "contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h"
 #include "core/providers/cuda/cuda_provider_options.h"
 #endif
 
@@ -42,6 +44,33 @@ 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 expected;
+  };
+  const Case cases[] = {
+      {10240, 1, 80, 48, 12, 32},
+      {6144, 8, 80, 48, 12, 32},
+      {1024, 4, 80, 48, 12, 16},
+      {10240, 16, 80, 48, 12, 16},
+      {10240, 1, 8, 48, 12, 8},
+      {10240, 1, 80, 132, 12, 8},
+      {10240, 1, 80, 48, 9, 8},
+  };
+
+  for (const Case& c : cases) {
+    SCOPED_TRACE("N = " + std::to_string(c.n) + ", M = " + std::to_string(c.m));
+    EXPECT_EQ(onnxruntime::contrib::cuda::PickFp8MmaKSplit(
+                  c.n, c.m, c.windows, c.sm_count, c.compute_capability_major),
+              c.expected);
+  }
+}
+
 // 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) {
@@ -448,6 +477,105 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesBf16) {
   }
 }
 
+// Selection boundaries for the residency-hinted entry point, at a fixed device size so the
+// expectations do not move with the test machine.
+TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCorePinnedResidencyBoundaries) {
+  constexpr int sm_count = 132;
+  using onnxruntime::contrib::cuda::Fp8MmaGemvPinsResidency;
+
+  // ceil(N / 16) has to land in (2 * sm_count, 3 * sm_count] == (264, 396].
+  EXPECT_FALSE(Fp8MmaGemvPinsResidency(16 * 264, 16, 1, sm_count));
+  EXPECT_TRUE(Fp8MmaGemvPinsResidency(16 * 264 + 1, 16, 1, sm_count));
+  EXPECT_TRUE(Fp8MmaGemvPinsResidency(16 * 396, 16, 1, sm_count));
+  EXPECT_FALSE(Fp8MmaGemvPinsResidency(16 * 396 + 1, 16, 1, sm_count));
+  // 8-warp blocks regress under any explicit bounds, 32-warp blocks cannot host 3 blocks per SM,
+  // and 2 or 4 row tiles spill at the register cap that 3 resident blocks imply.
+  EXPECT_FALSE(Fp8MmaGemvPinsResidency(16 * 300, 8, 1, sm_count));
+  EXPECT_FALSE(Fp8MmaGemvPinsResidency(16 * 300, 32, 1, sm_count));
+  EXPECT_FALSE(Fp8MmaGemvPinsResidency(16 * 300, 16, 2, sm_count));
+  EXPECT_FALSE(Fp8MmaGemvPinsResidency(16 * 300, 16, 4, sm_count));
+}
+
+// Runs the residency-hinted kernel. It is a second instantiation of the same body, so what is
+// under test is the dispatch: nothing above reaches it, because which N selects it depends on the
+// device's SM count (N = 4098 is 257 column blocks, already one wave on anything from 86 SMs up).
+TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCorePinnedResidencyFp16) {
+  if (!HasCudaEnvironment(800)) {
+    GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight.";
+  }
+
+  cudaDeviceProp device_prop{};
+  int device_id = 0;
+  ASSERT_EQ(cudaGetDevice(&device_id), cudaSuccess);
+  ASSERT_EQ(cudaGetDeviceProperties(&device_prop, device_id), cudaSuccess);
+  const int sm_count = device_prop.multiProcessorCount;
+
+  constexpr int64_t k = 1024;  // 16 K windows, so KSplit stays at its full 16
+  constexpr int64_t block_size = 256;
+  constexpr int64_t k_blocks = k / block_size;
+  // Narrowest N above 2 blocks per SM. Past N = 8192 the launcher drops to 8 warps per block and
+  // stops hinting at all, so a device that large has no shape to test here.
+  const int64_t n_pinned = 16 * (2 * sm_count + 1);
+  if (n_pinned >= 8192) {
+    GTEST_SKIP() << "Device has " << sm_count << " SMs; the hinted window is above N = 8192.";
+  }
+
+  static const float kWeightValues[] = {1.0f, 2.0f, -1.0f};      // exact in E4M3
+  static const float kActValues[] = {1.0f, -1.0f, 0.5f, -0.5f};  // exact in FP16
+  // A ragged width in the same window leaves the last 16-column tile partly out of range.
+  for (const int64_t n : {n_pinned, n_pinned + 5}) {
+    const int k_split = onnxruntime::contrib::cuda::PickFp8MmaKSplit(
+        static_cast<int>(n), 1, static_cast<int>(k / 64), sm_count, device_prop.major);
+    ASSERT_TRUE(onnxruntime::contrib::cuda::Fp8MmaGemvPinsResidency(static_cast<int>(n), k_split, 1, sm_count))
+        << "N = " << n << " should take the hinted entry point on this device";
+
+    std::vector<Float8E4M3FN> b(static_cast<size_t>(n * k));
+    std::vector<float> b_scale(static_cast<size_t>(n * k_blocks));
+    for (int64_t col = 0; col < n; ++col) {
+      for (int64_t i = 0; i < k; ++i) {
+        b[static_cast<size_t>(col * k + i)] = Float8E4M3FN(kWeightValues[(col + i) % 3]);
+      }
+      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;
+      }
+    }
+
+    // Only one row tile is hinted, so M stops at 8.
+    for (const int64_t m : {1, 3, 8}) {
+      SCOPED_TRACE("N = " + std::to_string(n) + ", M = " + std::to_string(m));
+      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)] * kWeightValues[(col + i) % 3] *
+                   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);
+    }
+  }
+}
+
 // Lane-ownership probe for the tensor-core path.
 //
 // The tests above sum over the whole K axis, so a wrong lane -> (row, column) mapping could in