From ff87616399e17869a8872ded499e5d3b7f7ff696 Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Wed, 2 Sep 2026 18:28:55 -0700 Subject: [PATCH 1/4] Add FP8 GEMV KSplit32 scheduling for client SM12x GPUs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/math/matmul_block_scaled_fp8.cu | 11 ++++---- .../math/matmul_block_scaled_fp8_tiling.h | 25 +++++++++++++++++ .../matmul_block_scaled_fp8_test.cc | 28 +++++++++++++++++++ 3 files changed, 59 insertions(+), 5 deletions(-) create mode 100644 onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu index ac0ea06d6bc96..6de21c1f32f0a 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/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 #include @@ -905,10 +906,8 @@ 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; - } + const int k_split = PickFp8MmaKSplit( + n, m, windows, device_prop.multiProcessorCount, device_prop.major); const int mtiles = (m > 16) ? 4 : ((m > 8) ? 2 : 1); const dim3 mma_blocks{static_cast((n + 15) / 16)}; const auto launch_mma = [&]() { @@ -934,7 +933,9 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, launch_mma.template operator()(); } }; - 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>(); diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h new file mode 100644 index 0000000000000..d8167d93a98af --- /dev/null +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h @@ -0,0 +1,25 @@ +// 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; +} + +} // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc index bf7ac5f56baee..b59287a1df3d2 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -13,6 +13,7 @@ // evaluates to false and every test in this file is compiled out. #include +#include "contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h" #include "core/providers/cuda/cuda_provider_options.h" #endif @@ -42,6 +43,33 @@ std::vector MakeConstRowWeight(const std::vector& 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) { From acb4d27b946f7553709810f15e974272762ea405 Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Fri, 4 Sep 2026 02:39:03 -0700 Subject: [PATCH 2/4] Qualify FP8 KSplit32 for GB10 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/math/matmul_block_scaled_fp8.cu | 99 +++++++++--- .../math/matmul_block_scaled_fp8_tiling.h | 26 ++- .../matmul_block_scaled_fp8_test.cc | 148 ++++++++++++++++-- 3 files changed, 237 insertions(+), 36 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu index 6de21c1f32f0a..5d4c5e1c72e63 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu @@ -853,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("ORT_FP8_GEMV_KSPLIT", 0); + static int const match_n = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP8_GEMV_MATCH_N", 0); + static int const match_k = + onnxruntime::ParseEnvironmentVariableWithDefault("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("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(); @@ -885,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(y) + static_cast(kFp8MmaGemvTileM) * n * element_size, static_cast(a) + static_cast(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 @@ -906,8 +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; - const int k_split = PickFp8MmaKSplit( - n, m, windows, device_prop.multiProcessorCount, device_prop.major); + // 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((n + 15) / 16)}; const auto launch_mma = [&]() { @@ -934,7 +973,8 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, } }; if (k_split == 32) { - launch_for_ksplit.template operator()<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) { @@ -1024,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 diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h index d8167d93a98af..88fcf0f9e22df 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h @@ -5,18 +5,30 @@ namespace onnxruntime::contrib::cuda { -inline int PickFp8MmaKSplit(int n, int m, int windows, int sm_count, int compute_capability_major) { +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 kWideOutputMinN = 16384; + constexpr int kWideOutputMinWindows = 80; + constexpr int kLongReductionMinN = 5120; + constexpr int kLongReductionMinWindows = 128; - // 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; - } + // 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. + if (compute_capability_major == 12 && compute_capability_minor == 1 && + sm_count == 48 && m <= 8 && + ((n >= kWideOutputMinN && windows >= kWideOutputMinWindows) || + (n >= kLongReductionMinN && windows >= kLongReductionMinWindows))) { + k_split = 32; } return k_split; diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc index b59287a1df3d2..94f4c7e393308 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "gtest/gtest.h" #include "test/common/cuda_op_test_utils.h" #include "test/common/tensor_op_test_utils.h" @@ -8,6 +10,13 @@ #include "test/unittest_util/conversion.h" #include "test/util/include/scoped_env_vars.h" +#ifdef _WIN32 +#include +#else +#include +#include +#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. @@ -31,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(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(length) < path.size(), "readlink(/proc/self/exe) failed."); + path.resize(static_cast(length)); + return path; +#endif +} + // Builds a [N, K] FP8 E4M3 weight where every element of row r equals row_value[r]. std::vector MakeConstRowWeight(const std::vector& row_value, int64_t k) { std::vector b(static_cast(row_value.size()) * static_cast(k)); @@ -50,26 +75,131 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreKSplitSelection) { int windows; int sm_count; int compute_capability_major; + int compute_capability_minor; 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}, + {17408, 1, 80, 48, 12, 1, 32}, + {16384, 8, 80, 48, 12, 1, 32}, + {5120, 8, 128, 48, 12, 1, 32}, + {16383, 1, 80, 48, 12, 1, 8}, + {16384, 1, 79, 48, 12, 1, 8}, + {5119, 1, 128, 48, 12, 1, 16}, + {5120, 1, 127, 48, 12, 1, 16}, + {1024, 4, 80, 48, 12, 1, 16}, + {7168, 8, 80, 48, 12, 1, 16}, + {10240, 9, 80, 48, 12, 1, 8}, + {10240, 16, 80, 48, 12, 1, 8}, + {10240, 1, 8, 48, 12, 1, 8}, + {10240, 1, 80, 47, 12, 1, 8}, + {10240, 1, 80, 49, 12, 1, 8}, + {10240, 1, 80, 48, 12, 0, 8}, + {10240, 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)); + 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.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(); + ScopedEnvironmentVariables scoped_env_vars{EnvVarMap{ + {"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; + } + + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight."; + } + + 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 b(static_cast(n * k)); + std::vector b_ref(static_cast(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(col * k + i)] = Float8E4M3FN(value); + b_ref[static_cast(col * k + i)] = value; + } + } + std::vector b_scale(static_cast(n * k_blocks)); + for (int64_t col = 0; col < n; ++col) { + for (int64_t kb = 0; kb < k_blocks; ++kb) { + b_scale[static_cast(col * k_blocks + kb)] = + static_cast(1 + (col + kb) % 3) / 4.0f; + } + } + std::vector a(static_cast(m * k)); + for (int64_t row = 0; row < m; ++row) { + for (int64_t i = 0; i < k; ++i) { + a[static_cast(row * k + i)] = kActValues[(row + i) % 4]; + } + } + std::vector expected(static_cast(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(row * k + i)] * b_ref[static_cast(col * k + i)] * + b_scale[static_cast(col * k_blocks + i / block_size)]; + } + expected[static_cast(row * n + col)] = acc; + } + } + + { + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", block_size); + test.AddInput("A", {m, k}, FloatsToMLFloat16s(a)); + test.AddInput("B", {n, k}, b); + test.AddInput("b_scale", {n, k_blocks}, b_scale); + test.AddOutput("Y", {m, n}, FloatsToMLFloat16s(expected)); + test.SetOutputTolerance(0.005f); + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + { + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", block_size); + test.AddInput("A", {m, k}, FloatsToBFloat16s(a)); + test.AddInput("B", {n, k}, b); + test.AddInput("b_scale", {n, k_blocks}, b_scale); + test.AddOutput("Y", {m, n}, FloatsToBFloat16s(expected)); + test.SetOutputTolerance(0.05f); + std::vector> 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) { From da238410bb8906c43fa662ee18f980053f4b1335 Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Tue, 8 Sep 2026 13:08:06 -0700 Subject: [PATCH 3/4] Refine FP8 KSplit32 selection geometry Use FP8 GEMV output-block geometry for the SM121 thresholds and cover coherent boundaries plus measured very-wide output shapes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../math/matmul_block_scaled_fp8_tiling.h | 14 ++++++++--- .../matmul_block_scaled_fp8_test.cc | 25 ++++++++++++------- 2 files changed, 26 insertions(+), 13 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h index 88fcf0f9e22df..698db58b7f5c1 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8_tiling.h @@ -17,17 +17,23 @@ 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 kWideOutputMinN = 16384; + constexpr int kOutputColumnsPerBlock = 16; + constexpr int kWideOutputMinBlocks = 1024; + constexpr int kLongReductionMinBlocks = 320; constexpr int kWideOutputMinWindows = 80; - constexpr int kLongReductionMinN = 5120; 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 && - ((n >= kWideOutputMinN && windows >= kWideOutputMinWindows) || - (n >= kLongReductionMinN && windows >= kLongReductionMinWindows))) { + ((output_blocks >= kWideOutputMinBlocks && windows >= kWideOutputMinWindows) || + (output_blocks >= kLongReductionMinBlocks && windows >= kLongReductionMinWindows))) { k_split = 32; } diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc index 94f4c7e393308..ba23df4117760 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -82,19 +82,26 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreKSplitSelection) { {17408, 1, 80, 48, 12, 1, 32}, {16384, 8, 80, 48, 12, 1, 32}, {5120, 8, 128, 48, 12, 1, 32}, - {16383, 1, 80, 48, 12, 1, 8}, + {16369, 1, 80, 48, 12, 1, 32}, + {16368, 1, 80, 48, 12, 1, 8}, {16384, 1, 79, 48, 12, 1, 8}, - {5119, 1, 128, 48, 12, 1, 16}, + {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}, - {10240, 9, 80, 48, 12, 1, 8}, - {10240, 16, 80, 48, 12, 1, 8}, - {10240, 1, 8, 48, 12, 1, 8}, - {10240, 1, 80, 47, 12, 1, 8}, - {10240, 1, 80, 49, 12, 1, 8}, - {10240, 1, 80, 48, 12, 0, 8}, - {10240, 1, 80, 48, 9, 0, 8}, + {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) { From 57581eb30833a9a958b6dea152c728f07a5637f4 Mon Sep 17 00:00:00 2001 From: Baiju Meswani Date: Tue, 8 Sep 2026 14:29:15 -0700 Subject: [PATCH 4/4] Make forced FP8 KSplit32 test hermetic Force the tensor-core GEMV path in the child process and report unavailable CUDA environments as skipped in the parent test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test/contrib_ops/matmul_block_scaled_fp8_test.cc | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc index ba23df4117760..13c7fade47500 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -121,7 +121,13 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreKSplitSelection) { 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"}, @@ -135,10 +141,6 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreForcedKSplit32) { return; } - if (!HasCudaEnvironment(800)) { - GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight."; - } - constexpr int64_t m = 8; constexpr int64_t n = 17; constexpr int64_t k = 2112; // 33 windows exercise a ragged KSplit32 reduction.