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..7df2739c2231c 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 @@ -514,17 +515,17 @@ struct Fp8GemvMma<__nv_bfloat16> { }; template -__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; 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 +__global__ void MatMulBlockScaledFp8MmaGemvKernel(ORT_FP8_MMA_GEMV_PARAMS) { + Fp8MmaGemvBody(ORT_FP8_MMA_GEMV_ARGS); +} + +template +__global__ __launch_bounds__(32 * KSplit, 3) void MatMulBlockScaledFp8MmaGemvKernelPinned(ORT_FP8_MMA_GEMV_PARAMS) { + Fp8MmaGemvBody(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("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("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((n + 15) / 16)}; + const bool pin_residency = Fp8MmaGemvPinsResidency(n, k_split, mtiles, device_prop.multiProcessorCount); const auto launch_mma = [&]() { const dim3 mma_threads{32, KSplit}; - if (is_bf16) { - MatMulBlockScaledFp8MmaGemvKernel<<>>( - reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast(a), b, - weight_scale, reinterpret_cast(bias), act_scale, m, n, k, block_size, k_blocks); - } else { - MatMulBlockScaledFp8MmaGemvKernel<<>>( - reinterpret_cast(y), reinterpret_cast(a), b, - weight_scale, reinterpret_cast(bias), act_scale, m, n, k, block_size, k_blocks); +#define ORT_FP8_LAUNCH_MMA(kernel_name) \ + do { \ + if (is_bf16) { \ + kernel_name<<>>( \ + reinterpret_cast<__nv_bfloat16*>(y), reinterpret_cast(a), b, \ + weight_scale, reinterpret_cast(bias), act_scale, m, n, k, \ + block_size, k_blocks); \ + } else { \ + kernel_name<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), b, \ + weight_scale, reinterpret_cast(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()(); } }; - 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..ba613579b3fff --- /dev/null +++ b/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 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..b53b95ea28a2c 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/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 +#include +#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 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) { @@ -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(n), 1, static_cast(k / 64), sm_count, device_prop.major); + ASSERT_TRUE(onnxruntime::contrib::cuda::Fp8MmaGemvPinsResidency(static_cast(n), k_split, 1, sm_count)) + << "N = " << n << " should take the hinted entry point on this device"; + + std::vector b(static_cast(n * k)); + std::vector b_scale(static_cast(n * k_blocks)); + for (int64_t col = 0; col < n; ++col) { + for (int64_t i = 0; i < k; ++i) { + b[static_cast(col * k + i)] = Float8E4M3FN(kWeightValues[(col + i) % 3]); + } + 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; + } + } + + // 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 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)] * kWeightValues[(col + i) % 3] * + 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); + } + } +} + // 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