From 99c3d87eb0ec939ea3bfd40836dce4934ebd869f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 28 Jul 2026 07:23:49 +0000 Subject: [PATCH 1/9] fp8 blocked scaled gemm: fp32 widening for RowsPerWarp > 1 --- .../cuda/math/matmul_block_scaled_fp8.cu | 116 +++++++++++++++--- .../matmul_block_scaled_fp8_test.cc | 56 +++++++++ 2 files changed, 157 insertions(+), 15 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 e4d30cee1f3b2..df96d121e6c09 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu @@ -260,23 +260,87 @@ __global__ void MatMulBlockScaledFp8GemvKernel(AType* __restrict__ output, continue; } const int kb = koff[u] / block_size; + if constexpr (RowsPerWarp == 1) { + // One row: nothing to amortize, and the extra live fp32 values of the hoisted + // form measurably cost occupancy at the wide M == 1 tiles (<1,4,2>). #pragma unroll - for (int c = 0; c < ColsPerWarp; ++c) { - const int col = col_base + c; - if (col >= n) { - continue; + for (int c = 0; c < ColsPerWarp; ++c) { + const int col = col_base + c; + if (col >= n) { + continue; + } + __half2 b_half[8]; + ORT_FP8_GEMV_CVT16(b_raw[u][c], b_half); + const float b_scale = weight_scale[static_cast(col) * k_blocks + kb]; + float partial = 0.0f; + ORT_FP8_GEMV_DOT16(a_lo[u][0], a_hi[u][0], b_half, partial); + acc[0][c] += partial * b_scale; + } + } else { + // M > 1 (speculative decode / MTP verify): the naive form widens B to fp32 once + // per row and A once per column, so a lane spends 8 + 48 * RowsPerWarp + // instructions per 16 weight bytes and the kernel goes ALU-bound long before it + // saturates HBM (1.25 TB/s at M == 4 vs 2.35 TB/s at M == 1 on H200). Hoisting + // both widenings out of the inner loop drops that to + // 8 * C + 16 * R + 16 * C + 16 * R * C, i.e. 200 -> 120 per column at R=4, C=2. + // The fma order is unchanged, so the result is bit-identical to the scalar form. + // The 16-element chunk is consumed in two halves to keep only RowsPerWarp * 8 + // widened A values live at a time. + __half2 b_half[ColsPerWarp][8]; +#pragma unroll + for (int c = 0; c < ColsPerWarp; ++c) { + ORT_FP8_GEMV_CVT16(b_raw[u][c], b_half[c]); + } + float b_scale[ColsPerWarp]; +#pragma unroll + for (int c = 0; c < ColsPerWarp; ++c) { + const int col = col_base + c; + b_scale[c] = (col < n) ? weight_scale[static_cast(col) * k_blocks + kb] : 0.0f; + } + float partial[RowsPerWarp][ColsPerWarp] = {}; +#pragma unroll + for (int h = 0; h < 2; ++h) { + float a_f[RowsPerWarp][8]; +#pragma unroll + for (int r = 0; r < RowsPerWarp; ++r) { + const AVec2* av = reinterpret_cast(h == 0 ? &a_lo[u][r] : &a_hi[u][r]); +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 v = to_float2(av[i]); + a_f[r][2 * i] = v.x; + a_f[r][2 * i + 1] = v.y; + } + } +#pragma unroll + for (int c = 0; c < ColsPerWarp; ++c) { + float b_f[8]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const float2 v = __half22float2(b_half[c][4 * h + i]); + b_f[2 * i] = v.x; + b_f[2 * i + 1] = v.y; + } +#pragma unroll + for (int r = 0; r < RowsPerWarp; ++r) { +#pragma unroll + for (int j = 0; j < 8; ++j) { + partial[r][c] = fmaf(a_f[r][j], b_f[j], partial[r][c]); + } + } + } } - __half2 b_half[8]; - ORT_FP8_GEMV_CVT16(b_raw[u][c], b_half); - const float b_scale = weight_scale[static_cast(col) * k_blocks + kb]; #pragma unroll for (int r = 0; r < RowsPerWarp; ++r) { if (row_base + r >= m) { continue; } - float partial = 0.0f; - ORT_FP8_GEMV_DOT16(a_lo[u][r], a_hi[u][r], b_half, partial); - acc[r][c] += partial * b_scale; +#pragma unroll + for (int c = 0; c < ColsPerWarp; ++c) { + if (col_base + c >= n) { + continue; + } + acc[r][c] += partial[r][c] * b_scale[c]; + } } } } @@ -486,9 +550,19 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, // single wave, so widening each warp (ColsPerWarp) and pre-issuing loads (Unroll) buys // memory-level parallelism at no real occupancy cost. It only pays once N is large // enough that dividing the column count still leaves the GPU full -- below N == 4096 - // the wider tiles measured slower than one column per warp on H200, and for M > 1 the - // extra live registers per warp (RowsPerWarp * ColsPerWarp accumulators plus the - // pre-issued loads) cost more than the extra parallelism is worth. + // the wider tiles measured slower than one column per warp on H200. + // + // M in [2, 4] is speculative decode (the MTP verify forward is N+1 tokens wide). There + // RowsPerWarp > 1 already reads the weight once for every row, so the kernel is ALU + // bound rather than bandwidth bound, and ColsPerWarp pays for a second reason: it + // amortizes the fp32 widening of A across columns (see the hoisted path in the kernel). + // ColsPerWarp is capped so that gridDim.x stays >= ~128 blocks (H200 has 132 SMs) -- + // below that the lost column parallelism costs more than the saved instructions. + // Measured on H200 (us, M == 4, vs the previous and vs cuBLAS fp16): + // 8192x2048 10.9 cublas | 13.7 <4,1,1> | 9.7 <4,4,1> + // 4096x2048 8.3 cublas | 8.1 <4,1,1> | 6.9 <4,4,1> + // 2048x4096 8.3 cublas | 9.0 <4,1,1> | 7.9 <4,2,2> + // 512x2048 7.3 cublas | 5.0 <4,1,1> | 4.6 <4,1,2> if (m == 1) { if (n >= 8192) { launch.template operator()<1, 4, 2>(); @@ -498,9 +572,21 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, launch.template operator()<1, 1, 1>(); } } else if (m <= 2) { - launch.template operator()<2, 1, 1>(); + if (n >= 8192) { + launch.template operator()<2, 4, 1>(); + } else if (n >= 2048) { + launch.template operator()<2, 2, 1>(); + } else { + launch.template operator()<2, 1, 2>(); + } } else if (m <= 4) { - launch.template operator()<4, 1, 1>(); + if (n >= 4096) { + launch.template operator()<4, 4, 1>(); + } else if (n >= 2048) { + launch.template operator()<4, 2, 2>(); + } else { + launch.template operator()<4, 1, 2>(); + } } else { launch.template operator()<8, 1, 1>(); } 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 3dbfc6d88ccfb..bca9ce2212154 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -241,6 +241,62 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvDecodeWideTilesFp16) { } } +// Covers the M > 1 (speculative decode / MTP verify) GEMV dispatch. There RowsPerWarp is 2 or 4 +// and ColsPerWarp is chosen from N (1 / 2 / 4), which selects the hoisted-widening code path in +// the kernel. Every (RowsPerWarp, ColsPerWarp, Unroll) combination is exercised with a ragged N so +// the per-column bounds predication is hit, and both the row and the column vary so a mis-mapped +// row or column cannot pass by accident. +TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvSpeculativeDecodeTilesFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight."; + } + + // m = 2 -> RowsPerWarp 2, m = 3/4 -> RowsPerWarp 4 (m = 3 also leaves a ragged row tail). + // n picks ColsPerWarp 1 (n < 2048), 2 and 4; all leave a ragged column tail. + for (const int64_t m : {2, 3, 4}) { + for (const int64_t n : {1026, 2050, 4098, 8194}) { + constexpr int64_t k = 64; // K % 16 == 0 -> GEMV path; two 32-element K blocks + constexpr int64_t block_size = 32; + constexpr int64_t k_blocks = k / block_size; + + static const float kRowValues[] = {1.0f, 2.0f, -1.0f, -2.0f}; + std::vector row_value(static_cast(n)); + for (int64_t r = 0; r < n; ++r) { + row_value[static_cast(r)] = kRowValues[r % 4]; + } + std::vector b = MakeConstRowWeight(row_value, k); + std::vector b_scale(static_cast(n * k_blocks), 1.0f); + + // A[row, :] = row + 1 -> Y[row, col] = row_value[col] * k * (row + 1). + 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)] = static_cast(row + 1); + } + } + std::vector expected(static_cast(m * n)); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < n; ++col) { + expected[static_cast(row * n + col)] = + row_value[static_cast(col)] * static_cast(k) * static_cast(row + 1); + } + } + + 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.5f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + } +} + #endif // USE_CUDA && !DISABLE_FLOAT8_TYPES && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 } // namespace onnxruntime::test From 6f23fcd869df2710e97881e6bd666ef7663f1964 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 28 Jul 2026 07:43:10 +0000 Subject: [PATCH 2/9] prmt quad decode --- .../cuda/math/matmul_block_scaled_fp4.cu | 135 ++++++++++++++---- 1 file changed, 106 insertions(+), 29 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu index 6fcf2af957bb2..040bd9d6351ae 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -86,12 +86,22 @@ __global__ void AddBiasKernel(T* __restrict__ y, const T* __restrict__ bias, int // // __nv_cvt_fp4x2_to_halfraw2() is emulated in software on pre-Blackwell parts and // costs ~10 ALU ops per pair, which dominates the decode GEMV below (measured 3.5x -// slowdown on H200). Instead build the target float directly from the code bits: -// for c = s e1 e0 m the pattern (s << 15) | ((c & 7) << ) has -// exponent field == e and mantissa == m/2, i.e. exactly value * 2^-(bias - 1). -// Multiplying by 2^(bias - 1) afterwards recovers the value, and the e == 0 -// subnormal encodings fall out correctly as well. Verified exhaustively against the -// intrinsic for all 256 packed byte values, for both half and bfloat16. +// slowdown on H200). E2M1 has only eight magnitudes {0, 0.5, 1, 1.5, 2, 3, 4, 6}, +// so the 16-bit float bit pattern is built directly with a `prmt.b32` byte-select +// from packed magnitude constants plus a shifted sign bit. One prmt performs *four* +// magnitude lookups at once, so a whole 32-bit weight word (eight codes) decodes in +// ~14 instructions instead of ~48 for the per-byte bit-twiddling it replaces. This +// matters because the GEMV is instruction-issue bound, not bandwidth bound (it runs +// at 1.57 TB/s = 33% of H200 HBM peak). +// +// The magnitude bytes below are the exact half/bf16 encodings of the eight FP4 +// values, so the decoded value is bit-identical to the previous path (which produced +// value * 2^-(bias - 1) and multiplied by 2^(bias - 1) afterwards) and to the +// __nv_cvt_fp4x2_to_halfraw2() intrinsic. The scalar fallback keeps the host-side +// build working and is used when inline PTX is unavailable. +// +// Shared with the QMoE FP4 GEMV; see Fp4I2FConverter in +// contrib_ops/cuda/llm/fpA_intB_gemv/details.h for the same decode. template struct Fp4Cvt; @@ -99,12 +109,47 @@ template <> struct Fp4Cvt { using Traits = Vec2Traits; using T2 = typename Traits::Type2; - static __device__ __forceinline__ T2 Raw(uint32_t b) { - const uint32_t lo = ((b & 0x07u) << 9) | ((b & 0x08u) << 12); - const uint32_t hi = ((b & 0x70u) << 5) | ((b & 0x80u) << 8); - return bit_cast(lo | (hi << 16)); + + // Decodes four consecutive E2M1 codes. `mag_sel` holds the four 3-bit magnitudes as the four + // low nibbles (bit 3 cleared so prmt stays in byte-select mode rather than sign-replicate + // mode); `sgn_sel` holds the four sign bits as 0/1 nibbles. + static __device__ __forceinline__ void DecodeQuad(uint32_t mag_sel, uint32_t sgn_sel, + T2& lo2, T2& hi2) { +#if defined(__CUDA_ARCH__) + // Sign bytes: nibble 0 -> 0x00, nibble 1 -> 0x80 (bit 7 of the half high byte). + uint32_t sb; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(sb) : "r"(0x00008000u), "r"(0u), "r"(sgn_sel)); + // half high byte per magnitude (low byte is always 0): + // codes 0..3 -> {0x00, 0x38, 0x3C, 0x3E}, codes 4..7 -> {0x40, 0x42, 0x44, 0x46}. + uint32_t hb; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hb) : "r"(0x3E3C3800u), "r"(0x46444240u), "r"(mag_sel)); + hb |= sb; + // Expand {b0,b1,b2,b3} to {0,b0,0,b1} and {0,b2,0,b3}, pulling the zero low bytes from the + // second (all-zero) prmt operand. + uint32_t lo, hi; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(lo) : "r"(hb), "r"(0u), "n"(0x1404)); + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hi) : "r"(hb), "r"(0u), "n"(0x3424)); + lo2 = bit_cast(lo); + hi2 = bit_cast(hi); +#else + ScalarDecodeQuad(mag_sel, sgn_sel, lo2, hi2); +#endif + } + + static __device__ __forceinline__ void ScalarDecodeQuad(uint32_t mag_sel, uint32_t sgn_sel, + T2& lo2, T2& hi2) { + constexpr uint16_t kMag[8] = {0x0000u, 0x3800u, 0x3C00u, 0x3E00u, + 0x4000u, 0x4200u, 0x4400u, 0x4600u}; + uint16_t e[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + e[i] = static_cast(kMag[(mag_sel >> (4 * i)) & 0x7u] | + (((sgn_sel >> (4 * i)) & 0x1u) << 15)); + } + lo2 = bit_cast(static_cast(e[0]) | (static_cast(e[1]) << 16)); + hi2 = bit_cast(static_cast(e[2]) | (static_cast(e[3]) << 16)); } - static __device__ __forceinline__ T2 Scale() { return Traits::splat(16384.0f); } // 2^14 + static __device__ __forceinline__ T2 Mul(T2 a, T2 b) { return Traits::mul2(a, b); } static __device__ __forceinline__ float2 ToFloat2(T2 v) { return Traits::to_float2(v); } }; @@ -113,14 +158,43 @@ template <> struct Fp4Cvt { using Traits = Vec2Traits; using T2 = typename Traits::Type2; - static __device__ __forceinline__ T2 Raw(uint32_t b) { - const uint32_t lo = ((b & 0x07u) << 6) | ((b & 0x08u) << 12); - const uint32_t hi = ((b & 0x70u) << 2) | ((b & 0x80u) << 8); - return bit_cast(lo | (hi << 16)); + + static __device__ __forceinline__ void DecodeQuad(uint32_t mag_sel, uint32_t sgn_sel, + T2& lo2, T2& hi2) { +#if defined(__CUDA_ARCH__) + uint32_t sb; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(sb) : "r"(0x00008000u), "r"(0u), "r"(sgn_sel)); + // bf16 high byte {0x00,0x3F,0x3F,0x3F, 0x40,0x40,0x40,0x40} and + // low byte {0x00,0x00,0x80,0xC0, 0x00,0x40,0x80,0xC0}. + uint32_t hb, lb; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hb) : "r"(0x3F3F3F00u), "r"(0x40404040u), "r"(mag_sel)); + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(lb) : "r"(0xC0800000u), "r"(0xC0804000u), "r"(mag_sel)); + hb |= sb; + // bf16 needs both bytes: interleave low/high bytes of elements 0,1 and 2,3. + uint32_t lo, hi; + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(lo) : "r"(lb), "r"(hb), "n"(0x5140)); + asm("prmt.b32 %0, %1, %2, %3;" : "=r"(hi) : "r"(lb), "r"(hb), "n"(0x7362)); + lo2 = bit_cast(lo); + hi2 = bit_cast(hi); +#else + ScalarDecodeQuad(mag_sel, sgn_sel, lo2, hi2); +#endif } - static __device__ __forceinline__ T2 Scale() { - return Traits::splat(85070591730234615865843651857942052864.0f); // 2^126 + + static __device__ __forceinline__ void ScalarDecodeQuad(uint32_t mag_sel, uint32_t sgn_sel, + T2& lo2, T2& hi2) { + constexpr uint16_t kMag[8] = {0x0000u, 0x3F00u, 0x3F80u, 0x3FC0u, + 0x4000u, 0x4040u, 0x4080u, 0x40C0u}; + uint16_t e[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + e[i] = static_cast(kMag[(mag_sel >> (4 * i)) & 0x7u] | + (((sgn_sel >> (4 * i)) & 0x1u) << 15)); + } + lo2 = bit_cast(static_cast(e[0]) | (static_cast(e[1]) << 16)); + hi2 = bit_cast(static_cast(e[2]) | (static_cast(e[3]) << 16)); } + static __device__ __forceinline__ T2 Mul(T2 a, T2 b) { return Traits::mul2(a, b); } static __device__ __forceinline__ float2 ToFloat2(T2 v) { return Traits::to_float2(v); } }; @@ -150,7 +224,6 @@ __global__ void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, const uint8_t* b_row = b_packed + static_cast(col) * (k >> 1); const uint8_t* ws_row = weight_scale + static_cast(col) * k_blocks; - const T2 up = Cvt::Scale(); constexpr int kBlockSize = 16; constexpr int kElemsPerLane = 32; // two 16-element blocks const int stride = 32 * kElemsPerLane; // 1024 elements per warp iteration @@ -169,7 +242,7 @@ __global__ void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, // * Within a row, byte offsets are (koff / 2) for B and koff * sizeof(T) for A, and // koff is a multiple of kElemsPerLane == 32. const uint4 packed = *reinterpret_cast(b_row + (koff >> 1)); - const uint8_t* bytes = reinterpret_cast(&packed); + const uint32_t words[4] = {packed.x, packed.y, packed.z, packed.w}; const uint4* ap = reinterpret_cast(a_row + koff); uint4 a0 = ap[0], a1 = ap[1], a2 = ap[2], a3 = ap[3]; @@ -177,16 +250,20 @@ __global__ void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, reinterpret_cast(&a2), reinterpret_cast(&a3)}; // fp32 accumulation per 16-element scale block, matching the reference path. - float p0 = 0.0f; - float p1 = 0.0f; + // Each 32-bit weight word holds 8 codes = 4 T2 pairs; words 0/1 cover the first + // 16-element scale block and words 2/3 the second. + float p[2] = {0.0f, 0.0f}; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t mag = words[w] & 0x77777777u; + const uint32_t sgn = (words[w] >> 3) & 0x11111111u; + T2 b2[4]; + Cvt::DecodeQuad(mag, sgn, b2[0], b2[1]); + Cvt::DecodeQuad(mag >> 16, sgn >> 16, b2[2], b2[3]); #pragma unroll - for (int i = 0; i < 16; ++i) { - const T2 bb = Cvt::Mul(Cvt::Raw(bytes[i]), up); - const float2 pv = Cvt::ToFloat2(Cvt::Mul(av[i >> 2][i & 3], bb)); - if (i < 8) { - p0 += pv.x + pv.y; - } else { - p1 += pv.x + pv.y; + for (int j = 0; j < 4; ++j) { + const float2 pv = Cvt::ToFloat2(Cvt::Mul(av[w][j], b2[j])); + p[w >> 1] += pv.x + pv.y; } } @@ -194,7 +271,7 @@ __global__ void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, const int kb1 = kb0 + 1; const float s0 = e4m3_to_float(ws_row[kb0]); const float s1 = e4m3_to_float(ws_row[kb1]); - acc += p0 * s0 + p1 * s1; + acc += p[0] * s0 + p[1] * s1; } } From f888710ec992827c589369aae98b62ee85bcc9dd Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 28 Jul 2026 08:29:13 +0000 Subject: [PATCH 3/9] FP4 dense GEMV: grid-gated row tiling The decode GEMV mapped M onto gridDim.y, so the packed NVFP4 weight was streamed from DRAM once per row of A. Speculative decoding / MTP verify runs M = N_spec + 1 rows and paid M x the weight traffic. Template the kernel on RowsPerBlock (1, 2 or 4): a warp now accumulates several rows at once and the uint4 weight load, the prmt E2M1 decode and the two E4M3 block scales are shared across the tile. Per-row fp32 accumulation order is unchanged, so results are bit-identical to RowsPerBlock == 1. Tiling removes M from gridDim.y, so it is gated on the column grid ceil(N / 8) covering at least one full wave of SMs on its own. H200, M = 4, half: N = 248320 (lm_head) 31040 col blocks 615.8 -> 537.7 us (1.15x) N = 2048 (shared down) 256 col blocks 4.30 -> 3.33 us (1.29x) N = 512 (shared gate) 64 col blocks 3.47 -> 4.27 us (0.81x, gated off) Also pin __launch_bounds__ per tile. Left unconstrained nvcc compiles the 4-row tile to 66 registers, which drops the block from 4 to 3 per SM; the pinned budgets are spill-free and worth ~7% on the large shapes by themselves. Qwen3.6 NVFP4 MTP decode on H200: FP4 dense GEMV 1.099 -> 0.950 ms/step (1.16x), end-to-end 10.448 -> 10.306 ms/step (-1.4%). Set ORT_FP4_GEMV_ROW_TILING=0 to force RowsPerBlock == 1. --- .../cuda/matmul_block_scaled_fp4.md | 21 +- .../cuda/math/matmul_block_scaled_fp4.cu | 219 +++++++++++++----- .../matmul_block_scaled_fp4_test.cc | 119 ++++++++++ 3 files changed, 301 insertions(+), 58 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md index cb3ccde088f5d..0cb55195e7616 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md @@ -107,7 +107,7 @@ quantization, CUTLASS setup, and underutilized tensor-core GEMM work. - `block_size == 16`, - `K % 32 == 0`. -Each warp computes one output element `Y[row, col]`. A lane consumes 32 K +Each warp computes one output column `col`. A lane consumes 32 K elements per iteration, which is exactly two 16-element scale blocks. The kernel loads: @@ -118,6 +118,25 @@ loads: The per-block scales are folded into the partial sums and `weight_scale_2` is applied once after the warp reduction. Optional bias is fused in lane 0. +### Row tiling + +A warp produces `RowsPerBlock` rows of `Y` at once (1, 2 or 4). The packed weight +load and the E2M1 decode are shared by all rows in the tile, which matters for +speculative decoding / MTP verify where `M = N_spec + 1 > 1`. This trades grid +parallelism for reuse, because `M` no longer contributes to `gridDim.y`, so it is +only enabled when the column grid `ceil(N / 8)` covers at least one full wave of +SMs on its own. Measured on H200 (132 SMs, `M = 4`, FP16): + +| Shape | N | column blocks | `RowsPerBlock = 4` vs `1` | +| --- | --- | --- | --- | +| `lm_head` | 248320 | 31040 | 615.8 -> 537.7 us (1.15x) | +| shared `down_proj` | 2048 | 256 | 4.30 -> 3.33 us (1.29x) | +| shared `gate_up_proj` | 512 | 64 | 3.47 -> 4.27 us (0.81x) - gated off | + +Per-row fp32 accumulation order does not depend on `RowsPerBlock`, so results are +bit-identical across tilings. Set `ORT_FP4_GEMV_ROW_TILING=0` to force +`RowsPerBlock == 1`. + This kernel reads the original unswizzled `[N, K / 16]` scale layout. Experiments with the native SM120 swizzled scale layout for GEMV were slower; see [matmul_block_scaled_fp4_experiments.md](matmul_block_scaled_fp4_experiments.md). diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu index 040bd9d6351ae..afc2c81ed889c 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -12,6 +12,8 @@ #include #endif +#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" @@ -70,17 +72,42 @@ __global__ void AddBiasKernel(T* __restrict__ y, const T* __restrict__ bias, int y[idx] = from_float(to_float(y[idx]) + to_float(bias[col])); } +// Number of warps (= output columns) per thread block in the GEMV. +constexpr int kGemvWarpsPerBlock = 8; + +// Largest M tile the GEMV will fold into a single block. M is at most kGemvMaxM (8) at the +// call site; beyond 4 rows the register pressure from the per-row accumulators and the A +// fragments starts to cost more occupancy than the extra weight reuse buys. +constexpr int kGemvMaxRowsPerBlock = 4; + +// Occupancy target for the GEMV. Holding RowsPerBlock accumulators costs registers, and left +// unconstrained nvcc trades occupancy for scheduling freedom: the 4-row tile lands on 66 +// registers, which drops the block from 4 to 3 per SM and costs ~7% on the large shapes. Pin a +// register budget per tile instead; all three instantiations compile spill-free at these targets. +template +struct GemvMinBlocksPerSm { + static constexpr int value = (RowsPerBlock == 1) ? 6 : ((RowsPerBlock == 2) ? 5 : 4); +}; + // ----------------------------------------------------------------------------- // Fused NVFP4 weight-only GEMV fast path for the decode phase (small M). // -// Each warp computes one output element Y[row, col]. The 32 lanes cooperatively -// reduce over K reading the packed NVFP4 weight directly (two E2M1 values per -// byte) with 16-byte coalesced loads, so the weight is streamed exactly once and -// no [N, K] dequantized buffer is materialized. Each lane consumes 32 contiguous -// K elements = 16 packed bytes, which span exactly two 16-element blocks; the two -// per-block E4M3 scales are folded in per half. The global fp32 scale is applied -// once after the warp reduction. Runs on any architecture with NVFP4 conversion -// intrinsics (CUDA >= 12.8), including SM90 and SM120. +// Each warp computes RowsPerBlock output elements Y[row0 .. row0 + RowsPerBlock, col]. +// The 32 lanes cooperatively reduce over K reading the packed NVFP4 weight directly +// (two E2M1 values per byte) with 16-byte coalesced loads, so the weight is streamed +// exactly once per block and no [N, K] dequantized buffer is materialized. Each lane +// consumes 32 contiguous K elements = 16 packed bytes, which span exactly two +// 16-element blocks; the two per-block E4M3 scales are folded in per half. The global +// fp32 scale is applied once after the warp reduction. Runs on any architecture with +// NVFP4 conversion intrinsics (CUDA >= 12.8), including SM90 and SM120. +// +// RowsPerBlock > 1 amortizes the weight load and the E2M1 decode across several rows +// of A, which matters for speculative decoding / MTP verify where M = N_spec + 1 > 1. +// It trades grid parallelism for that reuse (M no longer contributes to gridDim.y), so +// the launcher only enables it when gridDim.x alone already fills the device; see +// Fp4GemvRowsPerBlock() below. The per-row fp32 accumulation order is independent of +// RowsPerBlock, so results are bit-identical across tilings. +// ----------------------------------------------------------------------------- // ----------------------------------------------------------------------------- // Fast NVFP4 (E2M1) -> half / bfloat16 conversion. // @@ -199,28 +226,37 @@ struct Fp4Cvt { static __device__ __forceinline__ float2 ToFloat2(T2 v) { return Traits::to_float2(v); } }; -template -__global__ void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, - const T* __restrict__ a, - const uint8_t* __restrict__ b_packed, - const uint8_t* __restrict__ weight_scale, - const float* __restrict__ weight_scale_2, - const T* __restrict__ bias, - int m, - int n, - int k, - int k_blocks) { +template +__global__ __launch_bounds__(32 * kGemvWarpsPerBlock, GemvMinBlocksPerSm::value) + void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, + const T* __restrict__ a, + const uint8_t* __restrict__ b_packed, + const uint8_t* __restrict__ weight_scale, + const float* __restrict__ weight_scale_2, + const T* __restrict__ bias, + int m, + int n, + int k, + int k_blocks) { using Cvt = Fp4Cvt; using T2 = typename Cvt::T2; const int lane = threadIdx.x; // 0..31 const int col = blockIdx.x * blockDim.y + threadIdx.y; // n - const int row = blockIdx.y; // m - if (row >= m || col >= n) { + const int row0 = blockIdx.y * RowsPerBlock; // m + if (row0 >= m || col >= n) { return; } - const T* a_row = a + static_cast(row) * k; + // Clamp the row base pointers instead of masking inside the K loop: a ragged tail tile + // recomputes the last row and discards it at the store, which keeps the inner loop + // branch-free. Only the last block of gridDim.y can be ragged, and only when + // RowsPerBlock does not divide M. + const T* a_rows[RowsPerBlock]; +#pragma unroll + for (int r = 0; r < RowsPerBlock; ++r) { + a_rows[r] = a + static_cast(row0 + r < m ? row0 + r : m - 1) * k; + } const uint8_t* b_row = b_packed + static_cast(col) * (k >> 1); const uint8_t* ws_row = weight_scale + static_cast(col) * k_blocks; @@ -228,7 +264,12 @@ __global__ void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, constexpr int kElemsPerLane = 32; // two 16-element blocks const int stride = 32 * kElemsPerLane; // 1024 elements per warp iteration - float acc = 0.0f; + float acc[RowsPerBlock]; +#pragma unroll + for (int r = 0; r < RowsPerBlock; ++r) { + acc[r] = 0.0f; + } + for (int base = 0; base < k; base += stride) { const int koff = base + lane * kElemsPerLane; if (koff < k) { @@ -244,50 +285,95 @@ __global__ void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, const uint4 packed = *reinterpret_cast(b_row + (koff >> 1)); const uint32_t words[4] = {packed.x, packed.y, packed.z, packed.w}; - const uint4* ap = reinterpret_cast(a_row + koff); - uint4 a0 = ap[0], a1 = ap[1], a2 = ap[2], a3 = ap[3]; - const T2* av[4] = {reinterpret_cast(&a0), reinterpret_cast(&a1), - reinterpret_cast(&a2), reinterpret_cast(&a3)}; - - // fp32 accumulation per 16-element scale block, matching the reference path. + // Decode the 32 weights owned by this lane once; every row of A reuses them. // Each 32-bit weight word holds 8 codes = 4 T2 pairs; words 0/1 cover the first // 16-element scale block and words 2/3 the second. - float p[2] = {0.0f, 0.0f}; + T2 b2[4][4]; #pragma unroll for (int w = 0; w < 4; ++w) { const uint32_t mag = words[w] & 0x77777777u; const uint32_t sgn = (words[w] >> 3) & 0x11111111u; - T2 b2[4]; - Cvt::DecodeQuad(mag, sgn, b2[0], b2[1]); - Cvt::DecodeQuad(mag >> 16, sgn >> 16, b2[2], b2[3]); -#pragma unroll - for (int j = 0; j < 4; ++j) { - const float2 pv = Cvt::ToFloat2(Cvt::Mul(av[w][j], b2[j])); - p[w >> 1] += pv.x + pv.y; - } + Cvt::DecodeQuad(mag, sgn, b2[w][0], b2[w][1]); + Cvt::DecodeQuad(mag >> 16, sgn >> 16, b2[w][2], b2[w][3]); } const int kb0 = koff / kBlockSize; const int kb1 = kb0 + 1; const float s0 = e4m3_to_float(ws_row[kb0]); const float s1 = e4m3_to_float(ws_row[kb1]); - acc += p[0] * s0 + p[1] * s1; + +#pragma unroll + for (int r = 0; r < RowsPerBlock; ++r) { + const uint4* ap = reinterpret_cast(a_rows[r] + koff); + + // fp32 accumulation per 16-element scale block, matching the reference path. + // One uint4 of A is live at a time so the 4-row tile stays inside its register budget. + float p[2] = {0.0f, 0.0f}; +#pragma unroll + for (int w = 0; w < 4; ++w) { + uint4 av4 = ap[w]; + const T2* av = reinterpret_cast(&av4); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 pv = Cvt::ToFloat2(Cvt::Mul(av[j], b2[w][j])); + p[w >> 1] += pv.x + pv.y; + } + } + acc[r] += p[0] * s0 + p[1] * s1; + } } } #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - acc += __shfl_down_sync(0xffffffffu, acc, offset); - } - if (lane == 0) { - float result = acc * (*weight_scale_2); - if (bias != nullptr) { - result += to_float(bias[col]); + for (int r = 0; r < RowsPerBlock; ++r) { + float v = acc[r]; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + v += __shfl_down_sync(0xffffffffu, v, offset); + } + if (lane == 0 && row0 + r < m) { + float result = v * (*weight_scale_2); + if (bias != nullptr) { + result += to_float(bias[col]); + } + y[static_cast(row0 + r) * n + col] = from_float(result); } - y[static_cast(row) * n + col] = from_float(result); } } +bool Fp4GemvRowTilingEnabled() { + static bool const enabled = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_ROW_TILING", 1) == 1; + return enabled; +} + +// Picks the M tile for the GEMV. +// +// Folding M into the block removes gridDim.y parallelism: the grid shrinks from +// ceil(N / warps) * M blocks to ceil(N / warps) * ceil(M / RowsPerBlock). That is a win when +// the kernel is weight-bound and the device is already saturated by the N dimension alone +// (the packed weight and its E2M1 decode are amortized over RowsPerBlock rows), and a loss +// when it leaves SMs idle. Measured on H200 (132 SMs, M = 4, half): +// +// N = 248320 (lm_head) 31040 col blocks 615.8 -> 537.7 us (1.15x) +// N = 2048 (shared down) 256 col blocks 4.30 -> 3.33 us (1.29x) +// N = 512 (shared gate) 64 col blocks 3.47 -> 4.27 us (0.81x, fewer blocks than SMs) +// +// So gate on the column grid covering at least one full wave of SMs on its own. +int Fp4GemvRowsPerBlock(int m, int n) { + if (m <= 1 || !Fp4GemvRowTilingEnabled()) { + return 1; + } + static int const sm_count = onnxruntime::llm::common::getMultiProcessorCount(); + const int col_blocks = (n + kGemvWarpsPerBlock - 1) / kGemvWarpsPerBlock; + if (col_blocks < sm_count) { + return 1; + } + // Only 1, 2 and 4 are instantiated: a 3-row tile spills and is no better than two 2-row + // blocks, and M > 4 splits across gridDim.y (M is at most 8 here). + return (m >= kGemvMaxRowsPerBlock) ? kGemvMaxRowsPerBlock : 2; +} + } // namespace #endif // CUDA_VERSION >= 12080 @@ -389,21 +475,40 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, ORT_RETURN_IF_NOT(block_size == 16, "MatMulBlockQuantizedFp4Weight GEMV requires block_size == 16, got ", block_size, "."); ORT_RETURN_IF_NOT(k % 32 == 0, "MatMulBlockQuantizedFp4Weight GEMV requires K divisible by 32, got ", k, "."); const int k_blocks = (k + block_size - 1) / block_size; - constexpr int kWarpsPerBlock = 8; - const dim3 threads{32, kWarpsPerBlock}; - const dim3 blocks{static_cast((n + kWarpsPerBlock - 1) / kWarpsPerBlock), - static_cast(m)}; + const int rows_per_block = Fp4GemvRowsPerBlock(m, n); + const dim3 threads{32, kGemvWarpsPerBlock}; + const dim3 blocks{static_cast((n + kGemvWarpsPerBlock - 1) / kGemvWarpsPerBlock), + static_cast((m + rows_per_block - 1) / rows_per_block)}; const uint8_t* bp = reinterpret_cast(b_packed); const uint8_t* ws = reinterpret_cast(weight_scale); + +#define ORT_DISPATCH_FP4_GEMV(T) \ + do { \ + switch (rows_per_block) { \ + case 4: \ + MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ + reinterpret_cast(bias), m, n, k, k_blocks); \ + break; \ + case 2: \ + MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ + reinterpret_cast(bias), m, n, k, k_blocks); \ + break; \ + default: \ + MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ + reinterpret_cast(bias), m, n, k, k_blocks); \ + break; \ + } \ + } while (0) + if (is_bf16) { - MatMulBlockQuantizedFp4WeightGemvKernel<<>>( - reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, - reinterpret_cast(bias), m, n, k, k_blocks); + ORT_DISPATCH_FP4_GEMV(nv_bfloat16); } else { - MatMulBlockQuantizedFp4WeightGemvKernel<<>>( - reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, - reinterpret_cast(bias), m, n, k, k_blocks); + ORT_DISPATCH_FP4_GEMV(half); } +#undef ORT_DISPATCH_FP4_GEMV return CUDA_CALL(cudaGetLastError()); #else ORT_UNUSED_PARAMETER(y); diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc index a057e39dea2c6..32fa2db3e1045 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc @@ -329,6 +329,125 @@ TEST(MatMulBlockQuantizedFp4WeightOpTest, ZeroKWithBiasBf16) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +// Exercises the row-tiled GEMV (RowsPerBlock > 1). The tiling is only selected when the column +// grid (ceil(N / 8) blocks) already covers a full wave of SMs, so N must be large: N = 2048 gives +// 256 column blocks, which is above the SM count of every current device. Weights and per-block +// scales vary by column and A varies by row so that a mis-indexed row or column is caught. +// +// The M values sweep every tile the launcher can pick: 2 and 4 are exact tiles, 3 and 5 are ragged +// (gridDim.y == 2 with a partly out-of-range last block, whose extra rows are clamped and dropped). +// All expected values are multiples of K below 1024, hence exact in both FP16 and BF16. +namespace { + +constexpr int64_t kRowTiledN = 2048; + +// Weight row `col` is +1.0 (byte 0x22) or +2.0 (byte 0x44); its E4M3 block scale is 2.0 (0x40) +// on every third row and 1.0 (0x38) elsewhere. +float RowTiledWeightValue(int64_t col) { return (col % 2 == 0) ? 1.0f : 2.0f; } +float RowTiledScaleValue(int64_t col) { return (col % 3 == 0) ? 2.0f : 1.0f; } + +void MakeRowTiledWeights(int64_t n, int64_t k, int64_t k_blocks, + std::vector& b, std::vector& weight_scale) { + b.assign(n * (k / 2), 0); + weight_scale.assign(n * k_blocks, 0); + for (int64_t col = 0; col < n; ++col) { + const uint8_t byte = (RowTiledWeightValue(col) == 1.0f) ? 0x22 : 0x44; + for (int64_t j = 0; j < k / 2; ++j) { + b[col * (k / 2) + j] = byte; + } + const uint8_t s = (RowTiledScaleValue(col) == 2.0f) ? 0x40 : 0x38; + for (int64_t blk = 0; blk < k_blocks; ++blk) { + weight_scale[col * k_blocks + blk] = s; + } + } +} + +// A[row, :] = row + 1, so Y[row, col] = W[col] * S[col] * K * (row + 1). +void MakeRowTiledActivationsAndExpected(int64_t m, int64_t n, int64_t k, + std::vector& a, std::vector& expected) { + a.assign(m * k, 0.0f); + expected.assign(m * n, 0.0f); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < k; ++col) { + a[row * k + col] = static_cast(row + 1); + } + for (int64_t col = 0; col < n; ++col) { + expected[row * n + col] = RowTiledWeightValue(col) * RowTiledScaleValue(col) * + static_cast(k) * static_cast(row + 1); + } + } +} + +} // namespace + +TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvDecodeRowTiledFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t n = kRowTiledN; + constexpr int64_t k = 64; + constexpr int64_t k_blocks = k / 16; + + std::vector b, weight_scale; + MakeRowTiledWeights(n, k, k_blocks, b, weight_scale); + const std::vector weight_scale_2 = {1.0f}; + + for (int m_val : {2, 3, 4, 5}) { + const int64_t m = m_val; + SCOPED_TRACE("M = " + std::to_string(m)); + std::vector a, expected; + MakeRowTiledActivationsAndExpected(m, n, k, a, expected); + + OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", 16); + test.AddInput("A", {m, k}, FloatsToMLFloat16s(a)); + test.AddInput("B", {n, k / 2}, b); + test.AddInput("weight_scale", {n, k_blocks}, weight_scale); + test.AddInput("weight_scale_2", {1}, weight_scale_2); + test.AddOutput("Y", {m, n}, FloatsToMLFloat16s(expected)); + test.SetOutputTolerance(0.5f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + +TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvDecodeRowTiledBf16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t n = kRowTiledN; + constexpr int64_t k = 32; + constexpr int64_t k_blocks = k / 16; + + std::vector b, weight_scale; + MakeRowTiledWeights(n, k, k_blocks, b, weight_scale); + const std::vector weight_scale_2 = {1.0f}; + + for (int m_val : {2, 3, 4, 5}) { + const int64_t m = m_val; + SCOPED_TRACE("M = " + std::to_string(m)); + std::vector a, expected; + MakeRowTiledActivationsAndExpected(m, n, k, a, expected); + + OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", 16); + test.AddInput("A", {m, k}, FloatsToBFloat16s(a)); + test.AddInput("B", {n, k / 2}, b); + test.AddInput("weight_scale", {n, k_blocks}, weight_scale); + test.AddInput("weight_scale_2", {1}, weight_scale_2); + test.AddOutput("Y", {m, n}, FloatsToBFloat16s(expected)); + test.SetOutputTolerance(0.5f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + #endif // USE_CUDA && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 } // namespace onnxruntime::test From b535cfbeb42a5b313999f22ace10cc57e01d772a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 28 Jul 2026 22:13:59 +0000 Subject: [PATCH 4/9] Run FP8 weight-only decode GEMV on tensor cores The fused FP8 GEMV re-widens both operands to FP32 and does scalar FMAs, so at M > 1 it is issue bound rather than bandwidth bound: at M = 4 a lane runs ~240 instructions per 32 weight bytes and effective bandwidth drops from ~2.35 TB/s (M = 1) to ~1.25 TB/s. Add MatMulBlockScaledFp8MmaGemvKernel, which does the dot products with mma.m16n8k16 and FP32 accumulate. The weight feeds the mma A operand and the activation the mma B operand, so both fragments match the tensors' natural row-major layouts with no transpose or shared-memory staging; the mma "M" extent becomes the output column count (16 per warp) and the mma "N" extent becomes M. Fragment loads are coalesced by permuting the K axis - legal because K is a reduction axis and the same permutation is applied to both operands - which lets each lane load one contiguous uint4 of weight bytes per 64-element K window and feed four mma steps from it. KSplit warps per block take a strided share of the K windows and reduce through shared memory, restoring the memory-level parallelism lost by widening each warp to 16 columns. Used on SM80+ when K % 64 == 0, K >= 256 and block_size % 64 == 0; otherwise the existing kernel runs unchanged. Set ORT_FP8_GEMV_MMA=0 to force it. Not bit-identical to the FMA path but not less accurate: E4M3 to FP16/BF16 is lossless, the products are exact in FP32, and both kernels accumulate in FP32. Scored against an FP64 reference the maximum error is identical for the two. H200, standalone (us), FMA kernel -> mma kernel: 8192x2048 M=1 6.3 -> 5.1 M=4 9.8 -> 5.2 M=8 17.5 -> 5.7 4096x2048 M=1 4.8 -> 4.0 M=4 6.9 -> 4.1 M=8 10.1 -> 4.3 2048x4096 M=1 5.1 -> 4.2 M=4 8.0 -> 4.4 M=8 13.0 -> 4.5 512x2048 M=1 3.3 -> 3.1 M=4 4.7 -> 3.3 M=8 6.4 -> 3.3 Faster at every measured M, so it is preferred whenever its preconditions hold rather than only for M > 1. On a 40-layer Qwen3.6 NVFP4 MTP decode (130 FP8 matmul nodes per step, M = 4) the FP8 GEMV kernel family goes 1.052 -> 0.713 ms/step and the model goes 9.80 -> 9.54 ms/step, with no other kernel family affected. --- .../cuda/matmul_block_scaled_fp8.md | 27 ++ .../matmul_block_scaled_fp8_experiments.md | 103 ++++++- .../cuda/math/matmul_block_scaled_fp8.cc | 1 + .../cuda/math/matmul_block_scaled_fp8.cu | 269 +++++++++++++++++- .../cuda/math/matmul_block_scaled_fp8.h | 2 + .../matmul_block_scaled_fp8_test.cc | 140 +++++++++ 6 files changed, 537 insertions(+), 5 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md index c54fc3da51b9c..b83ed13bcc6fa 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md @@ -129,6 +129,33 @@ This path avoids a materialized dequant buffer and runs on all supported CUDA architectures because it uses regular FP8 conversion and warp-shuffle reduction, not architecture-specific block-scaled tensor cores. +### 4.1 Tensor-Core Sub-Path (SM80+) + +When the device is SM80 or newer and, in addition to the predicate above, + +- `K % 64 == 0` and `K >= 256`, +- `block_size % 64 == 0`, + +the GEMV runs `MatMulBlockScaledFp8MmaGemvKernel`, which replaces the FP32 FMA +dot products with `mma.m16n8k16` (FP32 accumulate). The FMA kernel is ALU bound +for `M > 1` because it re-widens `A` and `B` to FP32 for every row/column pair; +the tensor-core kernel cuts instructions per weight byte roughly 10x. + +The weight is fed to the mma **A** operand and the activation to the **B** +operand, so both fragments match the tensors' natural row-major layouts. The mma +"M" extent is therefore the output column count (16 columns per warp) and the mma +"N" extent is `M` (up to 8 rows). Fragment loads are made fully coalesced by +permuting the K axis - K is a reduction axis, so any permutation applied to both +operands leaves the result unchanged - which lets each lane load one contiguous +`uint4` of weight bytes per 64-element K window and feed four mma instructions +with it. `KSplit` warps per block take a strided share of the K windows and are +reduced through shared memory, which restores the memory-level parallelism lost +by giving each warp 16 columns instead of 1-4. + +Accuracy is unchanged: E4M3 to FP16/BF16 is lossless, the products are exact and +the mma accumulates in FP32 just like the FMA path. Set `ORT_FP8_GEMV_MMA=0` to +fall back to the FMA kernel. + --- ## 5. Default Path - Dequantize + cuBLAS diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md index a29fba69371ca..99b39053d635c 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md @@ -18,8 +18,9 @@ Related documentation: 3. [Prefill Bottleneck - Weight Dequantization](#3-prefill-bottleneck---weight-dequantization) 4. [Optimization - Vectorized Dequantization Kernel](#4-optimization---vectorized-dequantization-kernel) 5. [Decode GEMV - Memory-Level Parallelism](#5-decode-gemv---memory-level-parallelism) -6. [Benchmark Commands](#6-benchmark-commands) -7. [Lessons](#7-lessons) +6. [Decode GEMV - Tensor Cores](#6-decode-gemv---tensor-cores) +7. [Benchmark Commands](#7-benchmark-commands) +8. [Lessons](#8-lessons) --- @@ -253,7 +254,101 @@ preferred on speed alone. --- -## 6. Benchmark Commands +## 6. Decode GEMV - Tensor Cores + +Section 5 tuned memory-level parallelism at `M = 1`. At `M = 4` - the width of a +speculative-decode / MTP verify forward - the kernel is limited by something +else. With `RowsPerWarp = 4` a lane executes roughly 240 instructions per 32 +weight bytes, only 128 of which are the FMAs that do useful work, and effective +bandwidth falls from about 2.35 TB/s at `M = 1` to about 1.25 TB/s at `M = 4`. +More ILP cannot fix that; the dot products have to leave the FMA pipe. + +### 6.1 Design + +`MatMulBlockScaledFp8MmaGemvKernel` uses `mma.m16n8k16` with FP32 accumulation. +The operand assignment is the key decision: + +| mma operand | fed from | why it fits | +|---|---|---| +| `A[16, 16]` row-major | weight `[16 output cols][16 k]` | `B` is `[N, K]` row-major | +| `B[16, 8]` col-major | activation `[16 k][8 rows]` | `A` is `[M, K]` row-major | +| `D[16, 8]` | `y[16 output cols][8 rows]` | | + +So the mma "M" extent is the output column count and the mma "N" extent is `M`. +At `M = 4` half the mma N lanes are idle, which is irrelevant: the kernel is +bound by weight traffic and instruction issue, and both improve about 10x per +weight byte. + +The naive fragment load is badly coalesced - a lane needs bytes +`{2t, 2t+1, 2t+8, 2t+9}` of a row, which spreads a warp across 16 rows x 16 +bytes and over-fetches every 32-byte sector 2x. The fix is to **permute the K +axis**. K is a reduction axis, so any permutation applied to *both* operands +leaves the result unchanged. Inside a 64-element K window the permutation used is + +``` +mma k-slot (of step j) -> actual k + 2t, 2t+1 16t + 4j, 16t + 4j + 1 + 2t+8, 2t+9 16t + 4j + 2, 16t + 4j + 3 +``` + +so lane `(g = lane >> 2, t = lane & 3)` loads one contiguous `uint4` of weight +bytes `[16t, 16t + 16)` and the matching 32 activation bytes, four lanes cover 64 +contiguous bytes of one weight row, and that single `uint4` feeds all four mma +steps. + +16 columns per warp gives about 8x fewer warps than the FMA kernel, which alone +costs more in lost memory-level parallelism than the instruction saving is worth. +`KSplit` warps per block therefore take a strided share of the K windows and are +reduced through shared memory at the end. `KSplit = 8` for `N >= 8192` (the +column count already fills the grid) and 16 otherwise. + +Preconditions: SM80+, `K % 64 == 0`, `K >= 256`, `block_size % 64 == 0`, `M <= 8`. +Otherwise the FMA kernel runs unchanged. `ORT_FP8_GEMV_MMA=0` forces the FMA +kernel for A/B testing in a single binary. + +### 6.2 Accuracy + +Not bit-identical to the FMA kernel (different summation order), but not less +accurate either. E4M3 to FP16 is lossless, E4M3 to BF16 is lossless, FP16 x FP16 +products are exact in FP32, and the mma accumulates in FP32 exactly as the FMA +path does. Scored against an FP64 CPU reference on the shapes below, the maximum +error is *identical* for the two kernels (2-4e-4, i.e. pure FP16 output +rounding). + +### 6.3 Results (H200, us, standalone, `fp8_gemv_m4_bench.cu`) + +| Shape (N x K) | M | cuBLAS FP16 | FMA kernel | mma kernel | vs FMA | vs cuBLAS | +|---|---|---|---|---|---|---| +| 8192 x 2048 | 1 | 11.0 | 6.3 | **5.1** | 1.23x | 2.16x | +| 4096 x 2048 | 1 | 8.4 | 4.8 | **4.0** | 1.20x | 2.09x | +| 2048 x 4096 | 1 | 9.0 | 5.1 | **4.2** | 1.21x | 2.13x | +| 512 x 2048 | 1 | 6.7 | 3.3 | **3.1** | 1.06x | 2.12x | +| 8192 x 2048 | 4 | 11.0 | 9.8 | **5.2** | 1.87x | 2.10x | +| 4096 x 2048 | 4 | 8.5 | 6.9 | **4.1** | 1.69x | 2.09x | +| 2048 x 4096 | 4 | 8.5 | 8.0 | **4.4** | 1.82x | 1.92x | +| 512 x 2048 | 4 | 6.8 | 4.7 | **3.3** | 1.43x | 2.08x | +| 8192 x 2048 | 8 | 10.9 | 17.5 | **5.7** | 3.09x | 1.93x | +| 2048 x 4096 | 8 | 8.5 | 13.0 | **4.5** | 2.90x | 1.90x | + +The mma kernel is faster at every measured `M`, so it is preferred whenever its +preconditions hold rather than only for `M > 1`. Note also that the FMA kernel +crosses over and loses to cuBLAS at `M = 8`, while the mma kernel stays about 1.9x +ahead. + +End to end on a 40-layer Qwen3.6-35B-A3B NVFP4 MTP decode (130 FP8 matmul nodes +per step, `M = 4`), CUDA graphs on: + +| | FMA kernel | mma kernel | +|---|---|---| +| FP8 GEMV kernel time | 1.052 ms/step | **0.713 ms/step** | +| total kernel time | 7.368 ms/step | **7.021 ms/step** | +| wall | 9.80 ms/step | **9.54 ms/step** | + +No other kernel family moved. This optimization is kept. + +--- + +## 7. Benchmark Commands The commands below use `ORT_REPO` and `ORT_BUILD` so they can be copied without editing developer-specific paths. Set them once: @@ -307,7 +402,7 @@ CUDA_VISIBLE_DEVICES=0 "$ORT_BUILD/onnxruntime_provider_test" \ --- -## 7. Lessons +## 8. Lessons - The prefill path is memory bound on weight dequantization, not on the GEMM; latency there scales with `N*K` and is independent of `M`. diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc index 4c0c6d707f25c..172f534ba0ae8 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc @@ -119,6 +119,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons k_i, SafeInt(block_size_), std::is_same::value, + GetDeviceProp().major, Stream(context)); } 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 df96d121e6c09..1c81f67f9b49e 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu @@ -6,6 +6,7 @@ #include #include +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cu_inc/common.cuh" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" @@ -379,6 +380,237 @@ __global__ void MatMulBlockScaledFp8GemvKernel(AType* __restrict__ output, } } +// ----------------------------------------------------------------------------- +// Tensor-core decode GEMV (SM80+): mma.m16n8k16 with FP32 accumulation. +// +// The fp32-FMA kernel above is ALU bound once M > 1 (it re-widens A and B to fp32 for every +// row/column pair), so more ILP cannot help. The fix is to stop doing the dot products on the +// FMA pipe. Operand mapping -- the WEIGHT goes into the mma A slot and the ACTIVATION into the +// mma B slot, which makes both fragments match the tensors' natural layouts: +// +// mma A[16, 16] row-major <- weight[16 output columns][16 k] (B is [N, K] row-major) +// mma B[16, 8] col-major <- activation[16 k][8 rows] (A is [M, K] row-major) +// mma D[16, 8] -> y[16 output columns][8 rows] +// +// The mma "M" extent is therefore our N (16 output columns per warp) and the mma "N" extent is +// our M (up to 8 rows). At M == 4 half the mma N lanes are idle, but the instruction count per +// weight byte still drops about 10x, which is what the ALU-bound path needs. +// +// Fragment loads are made fully coalesced by PERMUTING the K axis. K is a reduction axis, so any +// permutation applied to BOTH operands leaves the result unchanged. Within a 64-element K window, +// lane (g = lane >> 2, t = lane & 3) loads one contiguous uint4 of weight bytes [16t, 16t + 16) +// and the matching 32 contiguous activation bytes; mma step j (0..3) then consumes bytes +// 4j..4j+3 of them. Four lanes cover 64 contiguous bytes of one weight row (no sector +// over-fetch) and a single uint4 load feeds four mma instructions. +// +// KSplit warps per block each take a strided share of the K windows and are reduced through +// shared memory. Without it, 16 columns per warp gives ~8x fewer warps than the fp32 kernel and +// the GPU runs out of memory-level parallelism long before it runs out of bandwidth. +// +// Requires k % 64 == 0 and block_size % 64 == 0 (so a 64-element window lies in one K block). +// Accuracy is unchanged: FP8 E4M3 -> FP16/BF16 is lossless, products are exact, and the mma +// accumulates in FP32 exactly like the FMA path. Measured against an FP64 reference on the Qwen +// shapes, max error is identical to the FMA kernel (~2-4e-4, i.e. pure FP16 output rounding). +template +struct Fp8GemvMma; + +template <> +struct Fp8GemvMma { + // 16 packed FP8 bytes -> 8 b32 registers, each holding 2 halves. + __device__ __forceinline__ static void Cvt16(const uint4& raw, uint32_t (&out)[8]) { + const __nv_fp8x2_storage_t* p = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const __half2 v = __half2(__nv_cvt_fp8x2_to_halfraw2(p[i], __NV_E4M3)); + out[i] = *reinterpret_cast(&v); + } + } + + __device__ __forceinline__ static void Mma(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +#else + ORT_UNUSED_PARAMETER(d); + ORT_UNUSED_PARAMETER(a); + ORT_UNUSED_PARAMETER(b); +#endif + } +}; + +template <> +struct Fp8GemvMma<__nv_bfloat16> { + // No direct FP8 -> BF16 intrinsic, so go through the (lossless) FP8 -> FP16 converter and then + // FP16 -> FP32 -> BF16. Every step is exact for E4M3: 3 mantissa bits fit in BF16's 7 and the + // E4M3 exponent range is a strict subset of BF16's. + __device__ __forceinline__ static void Cvt16(const uint4& raw, uint32_t (&out)[8]) { + const __nv_fp8x2_storage_t* p = reinterpret_cast(&raw); +#pragma unroll + for (int i = 0; i < 8; ++i) { + const __half2 h = __half2(__nv_cvt_fp8x2_to_halfraw2(p[i], __NV_E4M3)); + const __nv_bfloat162 v = __float22bfloat162_rn(__half22float2(h)); + out[i] = *reinterpret_cast(&v); + } + } + + __device__ __forceinline__ static void Mma(float (&d)[4], const uint32_t (&a)[4], const uint32_t (&b)[2]) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +#else + ORT_UNUSED_PARAMETER(d); + ORT_UNUSED_PARAMETER(a); + ORT_UNUSED_PARAMETER(b); +#endif + } +}; + +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, + int m, + int n, + int k, + int block_size, + int k_blocks) { + using Mma = Fp8GemvMma; + + const int lane = threadIdx.x; + const int warp = threadIdx.y; + const int g = lane >> 2; // mma group: output-column offset in the tile AND activation row + const int t = lane & 3; // mma thread-in-group: k sub-offset + const int windows = k >> 6; + + const int col_lo = blockIdx.x * 16 + g; + const int col_hi = col_lo + 8; + const bool lo_ok = col_lo < n; + const bool hi_ok = col_hi < n; + + const bool a_ok = g < m; + const size_t a_off = static_cast(g) * k + (t << 4); + + float acc[4] = {0.f, 0.f, 0.f, 0.f}; // scaled, summed over K blocks + float accb[4] = {0.f, 0.f, 0.f, 0.f}; // unscaled, within the current K block + int cur_kb = -1; + + for (int wi = warp; wi < windows; wi += KSplit) { + const int k0 = wi << 6; + const int kb = k0 / block_size; + if (kb != cur_kb) { + if (cur_kb >= 0) { + const float s_lo = lo_ok ? weight_scale[static_cast(col_lo) * k_blocks + cur_kb] : 0.f; + const float s_hi = hi_ok ? weight_scale[static_cast(col_hi) * k_blocks + cur_kb] : 0.f; + acc[0] += accb[0] * s_lo; + acc[1] += accb[1] * s_lo; + acc[2] += accb[2] * s_hi; + acc[3] += accb[3] * s_hi; + accb[0] = accb[1] = accb[2] = accb[3] = 0.f; + } + cur_kb = kb; + } + + // Activation: 32 contiguous bytes of row g (row >= m reads as zero). + uint4 a_raw[2]; + if (a_ok) { + const uint4* ap = reinterpret_cast(input_a + a_off + k0); + a_raw[0] = ap[0]; + a_raw[1] = ap[1]; + } else { + a_raw[0] = make_uint4(0, 0, 0, 0); + a_raw[1] = make_uint4(0, 0, 0, 0); + } + const uint32_t* av = reinterpret_cast(a_raw); + + // Weight: one uint4 per half of the 16-column tile. + const uint4 w_lo = lo_ok + ? *reinterpret_cast(input_b + static_cast(col_lo) * k + k0 + (t << 4)) + : make_uint4(0, 0, 0, 0); + const uint4 w_hi = hi_ok + ? *reinterpret_cast(input_b + static_cast(col_hi) * k + k0 + (t << 4)) + : make_uint4(0, 0, 0, 0); + + uint32_t b_lo[8], b_hi[8]; + Mma::Cvt16(w_lo, b_lo); + Mma::Cvt16(w_hi, b_hi); + +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t ra[4] = {b_lo[2 * j], b_hi[2 * j], b_lo[2 * j + 1], b_hi[2 * j + 1]}; + const uint32_t rb[2] = {av[2 * j], av[2 * j + 1]}; + Mma::Mma(accb, ra, rb); + } + } + + if (cur_kb >= 0) { + const float s_lo = lo_ok ? weight_scale[static_cast(col_lo) * k_blocks + cur_kb] : 0.f; + const float s_hi = hi_ok ? weight_scale[static_cast(col_hi) * k_blocks + cur_kb] : 0.f; + acc[0] += accb[0] * s_lo; + acc[1] += accb[1] * s_lo; + acc[2] += accb[2] * s_hi; + acc[3] += accb[3] * s_hi; + } + + if constexpr (KSplit > 1) { + __shared__ float red[KSplit * 32 * 4]; + float* slot = red + (warp * 32 + lane) * 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + slot[i] = acc[i]; + } + __syncthreads(); + if (warp != 0) { + return; + } +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc[i] = 0.f; + } + for (int ws = 0; ws < KSplit; ++ws) { + const float* p = red + (ws * 32 + lane) * 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc[i] += p[i]; + } + } + } + + // Lane (g, t) owns y[2t, col_lo], y[2t + 1, col_lo], y[2t, col_hi], y[2t + 1, col_hi]. + const int row = t << 1; + const float bias_lo = (bias != nullptr && lo_ok) ? to_float(bias[col_lo]) : 0.f; + const float bias_hi = (bias != nullptr && hi_ok) ? to_float(bias[col_hi]) : 0.f; + if (row < m) { + if (lo_ok) { + output[static_cast(row) * n + col_lo] = from_float(acc[0] + bias_lo); + } + if (hi_ok) { + output[static_cast(row) * n + col_hi] = from_float(acc[2] + bias_hi); + } + } + if (row + 1 < m) { + if (lo_ok) { + output[static_cast(row + 1) * n + col_lo] = from_float(acc[1] + bias_lo); + } + if (hi_ok) { + output[static_cast(row + 1) * n + col_hi] = from_float(acc[3] + bias_hi); + } + } +} + +// 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", 1) == 1; + return enabled; +} + } // namespace #endif // !DISABLE_FLOAT8_TYPES && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 @@ -516,6 +748,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, int k, int block_size, bool is_bf16, + int sm_major, cudaStream_t stream) { #if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 if (m <= 0 || n <= 0 || k <= 0) { @@ -528,9 +761,42 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, ORT_RETURN_IF_NOT(block_size % 16 == 0, "MatMulBlockQuantizedFp8Weight GEMV requires block_size divisible by 16, got ", block_size, "."); const int k_blocks = (k + block_size - 1) / block_size; + const auto* b = reinterpret_cast(b_fp8); + + // Tensor-core path (SM80+). Beats the FMA kernel at every M on H200: 1.06-1.23x at M == 1 and + // 1.4-1.87x at M == 4, where the FMA kernel is ALU bound. Needs 64-element K windows, and at + // least 4 of them so KSplit warps have something to do. + if (sm_major >= 8 && 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 dim3 mma_blocks{static_cast((n + 15) / 16)}; + 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), m, n, k, block_size, k_blocks); + } else { + MatMulBlockScaledFp8MmaGemvKernel<<>>( + reinterpret_cast(y), reinterpret_cast(a), b, + weight_scale, reinterpret_cast(bias), m, n, k, block_size, k_blocks); + } + }; + if (k_split == 16) { + launch_mma.template operator()<16>(); + } else if (k_split == 8) { + launch_mma.template operator()<8>(); + } else { + launch_mma.template operator()<4>(); + } + return CUDA_CALL(cudaGetLastError()); + } + constexpr int kWarpsPerBlock = 8; const dim3 threads{32, kWarpsPerBlock}; - const auto* b = reinterpret_cast(b_fp8); const auto launch = [&]() { const int cols_per_block = kWarpsPerBlock * ColsPerWarp; const dim3 blocks{static_cast((n + cols_per_block - 1) / cols_per_block), @@ -602,6 +868,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, ORT_UNUSED_PARAMETER(k); ORT_UNUSED_PARAMETER(block_size); ORT_UNUSED_PARAMETER(is_bf16); + ORT_UNUSED_PARAMETER(sm_major); ORT_UNUSED_PARAMETER(stream); return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MatMulBlockQuantizedFp8Weight requires CUDA 11.8 or later."); #endif diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h index 605e3ef552d92..bba4c7ec9adc7 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h @@ -67,6 +67,7 @@ Status LaunchQuantizeDequantizeActivationFp8(void* a_out, // directly (no [N, K] dequant buffer). a is [M, K] activation (FP16/BF16), b_fp8 is [N, K] // FP8 E4M3, weight_scale is [N, ceil(K/block_size)] fp32, bias is an optional [N] vector (may be // null). Output y is [M, N] in the activation type. Requires k % 16 == 0 and block_size % 16 == 0. +// sm_major selects the tensor-core (mma.m16n8k16) variant, which needs SM80+. // Runs on any architecture with FP8 conversion intrinsics (CUDA >= 11.8). Status LaunchMatMulBlockScaledFp8Gemv(void* y, const void* a, @@ -78,6 +79,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, int k, int block_size, bool is_bf16, + int sm_major, cudaStream_t stream); } // 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 bca9ce2212154..fb0842985bd3e 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -297,6 +297,146 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvSpeculativeDecodeTilesFp16) { } } +// Covers the SM80+ tensor-core (mma.m16n8k16) GEMV dispatch, taken when K is a multiple of 64 +// with at least 4 K windows and block_size is a multiple of 64. It shares no code with the FMA +// kernel: the K axis is permuted, one warp owns a 16-column tile (each lane holding two columns +// eight apart and two rows), and the K split is reduced across warps through shared memory. So +// it needs its own coverage - ragged N and M, several K blocks with per-block scales that vary by +// column, and inputs that vary along every one of M, N and K so a mis-mapped index cannot pass. +TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight."; + } + + struct Case { + int64_t n, k, block_size; + }; + // k = 256/320 give 4/5 K windows, which selects the small-KSplit fallback; k = 1024 gives the + // full KSplit (16 below N = 8192, 8 at or above it). Every N is ragged modulo the 16-column tile. + const Case cases[] = {{18, 256, 64}, {1026, 320, 64}, {4098, 1024, 256}, {8194, 1024, 512}}; + 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 + for (const Case& c : cases) { + const int64_t k_blocks = c.k / c.block_size; + for (const int64_t m : {1, 3, 4, 8}) { + // Periods 3 (weight) and 4 (activation) are coprime, so no (row, col) pair sums to zero by + // symmetry. Scales are 1/256-based so the reference stays small enough to be exact in FP16. + std::vector b(static_cast(c.n * c.k)); + std::vector b_ref(static_cast(c.n * c.k)); + for (int64_t col = 0; col < c.n; ++col) { + for (int64_t i = 0; i < c.k; ++i) { + const float v = kWeightValues[(col + i) % 3]; + b[static_cast(col * c.k + i)] = Float8E4M3FN(v); + b_ref[static_cast(col * c.k + i)] = v; + } + } + std::vector b_scale(static_cast(c.n * k_blocks)); + for (int64_t col = 0; col < c.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) / 256.0f; + } + } + std::vector a(static_cast(m * c.k)); + for (int64_t row = 0; row < m; ++row) { + for (int64_t i = 0; i < c.k; ++i) { + a[static_cast(row * c.k + i)] = kActValues[(row + i) % 4]; + } + } + std::vector expected(static_cast(m * c.n)); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < c.n; ++col) { + float acc = 0.0f; + for (int64_t i = 0; i < c.k; ++i) { + acc += a[static_cast(row * c.k + i)] * b_ref[static_cast(col * c.k + i)] * + b_scale[static_cast(col * k_blocks + i / c.block_size)]; + } + expected[static_cast(row * c.n + col)] = acc; + } + } + + OpTester test("MatMulBlockQuantizedFp8Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", c.block_size); + test.AddInput("A", {m, c.k}, FloatsToMLFloat16s(a)); + test.AddInput("B", {c.n, c.k}, b); + test.AddInput("b_scale", {c.n, k_blocks}, b_scale); + test.AddOutput("Y", {m, c.n}, FloatsToMLFloat16s(expected)); + test.SetOutputTolerance(0.05f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + } +} + +// BF16 companion to GemvTensorCoreTilesFp16. BF16 has its own mma instruction and its own FP8 +// converter (E4M3 -> FP16 -> FP32 -> BF16), neither shared with the FP16 instantiation. Also +// covers the bias input on the tensor-core path. +TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesBf16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight."; + } + + constexpr int64_t n = 1026; // ragged modulo the 16-column tile + constexpr int64_t k = 256; + constexpr int64_t block_size = 64; + constexpr int64_t k_blocks = k / block_size; + constexpr int64_t m = 4; + 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 v = kWeightValues[(col + i) % 3]; + b[static_cast(col * k + i)] = Float8E4M3FN(v); + b_ref[static_cast(col * k + i)] = v; + } + } + 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) / 256.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 bias(static_cast(n)); + for (int64_t col = 0; col < n; ++col) { + bias[static_cast(col)] = static_cast(col % 5) - 2.0f; + } + 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 + bias[static_cast(col)]; + } + } + + 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.AddOptionalInputEdge(); // a_scale (skipped) + test.AddInput("bias", {n}, FloatsToBFloat16s(bias)); + test.AddOutput("Y", {m, n}, FloatsToBFloat16s(expected)); + test.SetOutputTolerance(0.1f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + #endif // USE_CUDA && !DISABLE_FLOAT8_TYPES && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 } // namespace onnxruntime::test From f545074b2f8a6f5afecca9b2f6953a70c38efdc1 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 29 Jul 2026 08:53:38 +0000 Subject: [PATCH 5/9] Run the NVFP4 decode GEMV on tensor cores The fused NVFP4 decode GEMV gave each warp one output column and reduced over K with __shfl_down. That makes the warp re-read the whole A tile once per output column: at RowsPerBlock = 4 it pulls 16 KB of activation for 1 KB of weight, a 16:1 amplification that dominated lm_head (N = 248320, 537.6 us on H200). On SM80+ with K % 128 == 0 and M <= 8, replace the shuffle reduction with mma.m16n8k16 so a warp produces 16 output columns at once. The fragments need no staging because the weight goes in the mma A slot and the activation in the mma B slot: B is [N, K] row-major, exactly the A-row-major fragment, and A is [M, K] row-major, exactly the B-col-major fragment. The K axis is then permuted so each lane's loads are contiguous, which is legal because K is a reduction axis and the same permutation applies to both operands. Since the mma sums across the four lanes holding different 16-element scale blocks, the accumulator cannot be flushed per block; the E4M3 scale is folded into the decoded weight before the mma instead. That is exact in FP16 and BF16 (E2M1 carries 2 significand bits and E4M3 carries 4, so the product needs at most 6) and cannot overflow (max 6 * 448 = 2688, min 0.5 * 2^-9 = 2^-10). KSplit warps per block take a strided share of the K windows and reduce through shared memory, without which the 16x drop in warp count costs more on the small MLP shapes than the extra reuse buys. Measured on H200 (132 SMs, M = 4, FP16), scalar -> tensor core: lm_head N = 248320, K = 2048 537.6 -> 108.3 us 4.96x gate_up_proj N = 512, K = 2048 3.48 -> 2.99 us 1.17x down_proj N = 2048, K = 512 3.32 -> 2.52 us 1.32x Over a Qwen3.6 NVFP4 MTP decode step this is 0.949 -> 0.448 ms/step for the FP4 GEMV family and 9.54 -> 8.99 ms/step end to end. The fp32 accumulation order differs from the scalar path, so results are not bit-identical to it, but max relative error against an fp64 reference is unchanged (2.4e-04 .. 3.5e-04, i.e. NVFP4 quantization noise). Set ORT_FP4_GEMV_MMA=0 to fall back. The new tests vary the weight, the block scales and the activation along K, so an inconsistent k mapping cannot pass; the existing GEMV tests use K = 32/64 and never reach this path. --- .../cuda/matmul_block_scaled_fp4.md | 54 +++ .../cuda/math/matmul_block_scaled_fp4.cc | 1 + .../cuda/math/matmul_block_scaled_fp4.cu | 342 ++++++++++++++++++ .../cuda/math/matmul_block_scaled_fp4.h | 3 + .../matmul_block_scaled_fp4_test.cc | 154 ++++++++ 5 files changed, 554 insertions(+) diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md index 0cb55195e7616..c629131c994a4 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md @@ -137,6 +137,58 @@ Per-row fp32 accumulation order does not depend on `RowsPerBlock`, so results ar bit-identical across tilings. Set `ORT_FP4_GEMV_ROW_TILING=0` to force `RowsPerBlock == 1`. +### Tensor-core sub-path (SM80+) + +On SM80 and newer, when `K % 128 == 0` and `M <= 8`, the warp reduction above is +replaced by `mma.m16n8k16`, so a warp produces **16 output columns** at once +instead of one. Set `ORT_FP4_GEMV_MMA=0` to fall back to the scalar path. + +The scalar path re-reads the whole `A` tile once per output column: at +`RowsPerBlock = 4` a warp pulls 16 KB of activation for 1 KB of weight, a 16:1 +amplification that dominates `lm_head`. Producing 16 columns per warp cuts those +re-reads by 16x. + +The fragments are free because the **weight goes in the mma `A` slot** and the +**activation in the mma `B` slot**: `B` is `[N, K]` row-major, which is exactly the +`A`-row-major fragment, and `A` is `[M, K]` row-major, which is exactly the +`B`-col-major fragment. No transpose, no `ldmatrix`, no shared-memory staging. The +mma `M` extent becomes the column count (16) and the mma `N` extent becomes `M`. + +The K axis is then permuted so the four k-slots a lane needs are contiguous in +memory, which is legal because K is a reduction axis and the same permutation is +applied to both operands. A window is 128 K elements = 64 packed bytes; lane +`(g, t)` owns elements `[32t, 32t + 32)`, i.e. one `uint4` of weight and one +`uint4` of activation, spanning exactly two 16-element scale blocks. + +Because the mma sums across the four `t` lanes, which hold *different* scale +blocks, the accumulator cannot be flushed per block. The E4M3 scale is instead +folded into the decoded weight before the mma. That is exact in both FP16 and +BF16: E2M1 magnitudes carry 2 significand bits and E4M3 scales carry 4, so the +product needs at most 6, inside FP16's 11 and BF16's 8; the range is safe too +(max `6 * 448 = 2688`, min `0.5 * 2^-9 = 2^-10`). + +`KSplit` warps per block take a strided share of the K windows and reduce through +shared memory. Without it, 16 columns per warp yields 16x fewer warps than the +scalar path and the small MLP shapes lose more to idle SMs than they gain. The +launcher picks `ColTiles = 4, KSplit = 2` when the column grid alone still covers +a wave of SMs, and otherwise `ColTiles = 1` with as many `KSplit` warps as there +are K windows. + +Measured on H200 (132 SMs, `M = 4`, FP16), scalar -> tensor core: + +| Shape | N | K | scalar | tensor core | speedup | +| --- | --- | --- | --- | --- | --- | +| `lm_head` | 248320 | 2048 | 537.6 us | 108.3 us | **4.96x** | +| shared `gate_up_proj` | 512 | 2048 | 3.48 us | 2.99 us | 1.17x | +| shared `down_proj` | 2048 | 512 | 3.32 us | 2.52 us | 1.32x | + +Over a Qwen3.6 NVFP4 MTP decode step (121 FP4 GEMV launches) this is +0.949 -> 0.448 ms/step, and 9.54 -> 8.99 ms/step end to end. + +The fp32 accumulation order differs from the scalar path, so this path is not +bit-identical to it. Max relative error against an fp64 reference is unchanged +(2.4e-04 .. 3.5e-04, i.e. NVFP4 quantization noise, not accumulation noise). + This kernel reads the original unswizzled `[N, K / 16]` scale layout. Experiments with the native SM120 swizzled scale layout for GEMV were slower; see [matmul_block_scaled_fp4_experiments.md](matmul_block_scaled_fp4_experiments.md). @@ -215,6 +267,8 @@ GEMM and the original scale tensor for the other paths. | Variable | Default | Meaning | |----------|---------|---------| | `ORT_MATMUL_BLOCK_SCALED_FP4_NATIVE_SM120` | `0` | Enables the opt-in native SM120 NVFP4 x NVFP4 GEMM path when the shape and device guards pass. | +| `ORT_FP4_GEMV_MMA` | `1` | Set to `0` to disable the decode GEMV tensor-core sub-path (`mma.m16n8k16`, SM80+, `K % 128 == 0`) and use the scalar warp-reduction path. | +| `ORT_FP4_GEMV_ROW_TILING` | `1` | Set to `0` to force `RowsPerBlock == 1` in the scalar decode GEMV. | The default remains the existing weight-only semantics: decode GEMV for small `M`, otherwise dequantize `B` and call cuBLAS. diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc index b0aa7d4dabe56..13b9b80dfd230 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc @@ -196,6 +196,7 @@ Status MatMulBlockQuantizedFp4Weight::ComputeImpl(OpKernelContext* context) cons k_i, SafeInt(block_size_), std::is_same::value, + sm_ / 10, Stream(context)); } diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu index afc2c81ed889c..10b0c80875bdd 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -6,6 +6,8 @@ #include #include #include +#include +#include #if defined(CUDA_VERSION) && CUDA_VERSION >= 12080 #include @@ -341,6 +343,248 @@ __global__ __launch_bounds__(32 * kGemvWarpsPerBlock, GemvMinBlocksPerSm> half's 2^-14 minimum). +// +// KSplit warps per block take a strided share of the K windows and reduce through shared +// memory. Without it, 16 columns per warp yields 16x fewer warps than the scalar path and +// the small MLP shapes lose more to idle SMs than they gain per warp. +// +// Note the fp32 accumulation order differs from the scalar path, so this path is not +// bit-identical to it. Measured max relative error against an fp64 reference is identical +// to the scalar path on every shape tested (2.4e-04 .. 3.5e-04, i.e. NVFP4 quantization +// noise, not accumulation noise). +template +struct Fp4GemvMma; + +template <> +struct Fp4GemvMma { + using T2 = half2; + + // E4M3 -> half is exact, and half broadcasts into both lanes of the half2 weight pair. + static __device__ __forceinline__ T2 BroadcastScale(uint8_t e4m3) { + const half h = static_cast( + __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(e4m3), __NV_E4M3)); + return __half2half2(h); + } + + static __device__ __forceinline__ uint32_t Pack(T2 v) { + uint32_t r; + memcpy(&r, &v, sizeof(r)); + return r; + } + + static __device__ __forceinline__ void Mma(float (&d)[4], const uint32_t (&a)[4], + const uint32_t (&b)[2]) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +#else + ORT_UNUSED_PARAMETER(d); + ORT_UNUSED_PARAMETER(a); + ORT_UNUSED_PARAMETER(b); +#endif + } +}; + +template <> +struct Fp4GemvMma { + using T2 = nv_bfloat162; + + // E4M3 has 4 significand bits and an exponent range of [-9, 8], so the half hop and the + // bfloat16 result are both exact. + static __device__ __forceinline__ T2 BroadcastScale(uint8_t e4m3) { + const half h = static_cast( + __nv_cvt_fp8_to_halfraw(static_cast<__nv_fp8_storage_t>(e4m3), __NV_E4M3)); + return __bfloat162bfloat162(__float2bfloat16(__half2float(h))); + } + + static __device__ __forceinline__ uint32_t Pack(T2 v) { + uint32_t r; + memcpy(&r, &v, sizeof(r)); + return r; + } + + static __device__ __forceinline__ void Mma(float (&d)[4], const uint32_t (&a)[4], + const uint32_t (&b)[2]) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + asm volatile( + "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " + "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};" + : "+f"(d[0]), "+f"(d[1]), "+f"(d[2]), "+f"(d[3]) + : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1])); +#else + ORT_UNUSED_PARAMETER(d); + ORT_UNUSED_PARAMETER(a); + ORT_UNUSED_PARAMETER(b); +#endif + } +}; + +// Decodes one 32-bit packed weight word (eight E2M1 codes) into four T2 pairs, so that +// out[i] holds elements {2i, 2i + 1}. +template +__device__ __forceinline__ void Fp4DecodeWord(uint32_t word, typename Fp4Cvt::T2 (&out)[4]) { + const uint32_t mag = word & 0x77777777u; + const uint32_t sgn = (word >> 3) & 0x11111111u; + Fp4Cvt::DecodeQuad(mag, sgn, out[0], out[1]); + Fp4Cvt::DecodeQuad(mag >> 16, sgn >> 16, out[2], out[3]); +} + +template +__global__ __launch_bounds__(32 * KSplit * ColTiles, 1) void MatMulBlockQuantizedFp4WeightMmaGemvKernel( + T* __restrict__ y, + const T* __restrict__ a, + const uint8_t* __restrict__ b_packed, + const uint8_t* __restrict__ weight_scale, + const float* __restrict__ weight_scale_2, + const T* __restrict__ bias, + int m, + int n, + int k, + int k_blocks) { + using Cvt = Fp4Cvt; + using MmaOp = Fp4GemvMma; + using T2 = typename Cvt::T2; + + const int lane = threadIdx.x; + const int g = lane >> 2; // mma A row -> output column within the 16-column tile + const int t = lane & 3; // mma k-slot -> which 32-element slice of the window + const int warp = static_cast(threadIdx.y); + const int warp_k = (KSplit == 1) ? 0 : (warp % KSplit); + const int warp_c = (KSplit == 1) ? warp : (warp / KSplit); + + const int col_lo = (static_cast(blockIdx.x) * ColTiles + warp_c) * 16 + g; + const int col_hi = col_lo + 8; + const bool lo_ok = col_lo < n; + const bool hi_ok = col_hi < n; + + // Out-of-range columns fold onto column 0 so every load stays in bounds; their results + // are masked off at the store. + const int half_k = k >> 1; + const uint8_t* b_lo = b_packed + static_cast(lo_ok ? col_lo : 0) * half_k; + const uint8_t* b_hi = b_packed + static_cast(hi_ok ? col_hi : 0) * half_k; + const uint8_t* ws_lo = weight_scale + static_cast(lo_ok ? col_lo : 0) * k_blocks; + const uint8_t* ws_hi = weight_scale + static_cast(hi_ok ? col_hi : 0) * k_blocks; + + const bool a_ok = g < m; + const T* a_row = a + static_cast(a_ok ? g : 0) * k; + + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + const int windows = k >> 7; + + for (int wi = warp_k; wi < windows; wi += KSplit) { + const int kbase = (wi << 7) + (t << 5); + const uint4 wl4 = *reinterpret_cast(b_lo + (kbase >> 1)); + const uint4 wh4 = *reinterpret_cast(b_hi + (kbase >> 1)); + const int kb = (wi << 3) + (t << 1); + const T2 sl[2] = {MmaOp::BroadcastScale(ws_lo[kb]), MmaOp::BroadcastScale(ws_lo[kb + 1])}; + const T2 sh[2] = {MmaOp::BroadcastScale(ws_hi[kb]), MmaOp::BroadcastScale(ws_hi[kb + 1])}; + const uint4* ap = reinterpret_cast(a_row + kbase); + const uint32_t wl[4] = {wl4.x, wl4.y, wl4.z, wl4.w}; + const uint32_t wh[4] = {wh4.x, wh4.y, wh4.z, wh4.w}; + +#pragma unroll + for (int w = 0; w < 4; ++w) { + T2 bl[4], bh[4]; + Fp4DecodeWord(wl[w], bl); + Fp4DecodeWord(wh[w], bh); + // Words 0,1 fall in the first 16-element scale block of this lane's slice, words 2,3 + // in the second. + const T2 s0 = sl[w >> 1]; + const T2 s1 = sh[w >> 1]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + bl[i] = Cvt::Mul(bl[i], s0); + bh[i] = Cvt::Mul(bh[i], s1); + } + const uint4 av4 = a_ok ? ap[w] : make_uint4(0u, 0u, 0u, 0u); + const uint32_t* av = reinterpret_cast(&av4); +#pragma unroll + for (int j = 0; j < 2; ++j) { + // A regs are (row g, k 2t..2t+1), (row g+8, ...), (row g, k 2t+8..), (row g+8, ...). + const uint32_t ra[4] = {MmaOp::Pack(bl[2 * j]), MmaOp::Pack(bh[2 * j]), + MmaOp::Pack(bl[2 * j + 1]), MmaOp::Pack(bh[2 * j + 1])}; + const uint32_t rb[2] = {av[2 * j], av[2 * j + 1]}; + MmaOp::Mma(acc, ra, rb); + } + } + } + + if constexpr (KSplit > 1) { + __shared__ float red[ColTiles * KSplit * 32 * 4]; + float* mine = red + ((warp_c * KSplit + warp_k) * 32 + lane) * 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + mine[i] = acc[i]; + } + __syncthreads(); + if (warp_k != 0) { + return; + } +#pragma unroll + for (int ws = 1; ws < KSplit; ++ws) { + const float* other = red + ((warp_c * KSplit + ws) * 32 + lane) * 4; +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc[i] += other[i]; + } + } + } + + // D regs are (row g, cols 2t, 2t+1) and (row g+8, cols 2t, 2t+1): mma rows are our output + // columns and mma columns are our M rows. + const float g2 = *weight_scale_2; + const int r0 = t << 1; + const float bias_lo = (bias != nullptr && lo_ok) ? to_float(bias[col_lo]) : 0.0f; + const float bias_hi = (bias != nullptr && hi_ok) ? to_float(bias[col_hi]) : 0.0f; + if (lo_ok) { + if (r0 < m) { + y[static_cast(r0) * n + col_lo] = from_float(acc[0] * g2 + bias_lo); + } + if (r0 + 1 < m) { + y[static_cast(r0 + 1) * n + col_lo] = from_float(acc[1] * g2 + bias_lo); + } + } + if (hi_ok) { + if (r0 < m) { + y[static_cast(r0) * n + col_hi] = from_float(acc[2] * g2 + bias_hi); + } + if (r0 + 1 < m) { + y[static_cast(r0 + 1) * n + col_hi] = from_float(acc[3] * g2 + bias_hi); + } + } +} + bool Fp4GemvRowTilingEnabled() { static bool const enabled = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_ROW_TILING", 1) == 1; @@ -374,6 +618,45 @@ int Fp4GemvRowsPerBlock(int m, int n) { return (m >= kGemvMaxRowsPerBlock) ? kGemvMaxRowsPerBlock : 2; } +bool Fp4GemvMmaEnabled() { + static bool const enabled = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_MMA", 1) == 1; + return enabled; +} + +struct Fp4MmaConfig { + int k_split; + int col_tiles; +}; + +// Picks the (KSplit, ColTiles) shape for the tensor-core GEMV. +// +// A warp owns 16 output columns, so the column grid is 16x smaller than the scalar path's. +// Wide column blocks (ColTiles = 4) give the best weight locality, but only pay off once N +// alone still covers a full wave of SMs; below that the block count collapses and the +// device idles, so columns are traded back for K-split warps. Measured on H200 (132 SMs, +// M = 4, half), scalar path -> best tensor-core config: +// +// N = 248320, K = 2048 (lm_head) 538.5 -> 107.7 us (5.00x) KSplit 2, ColTiles 4 +// N = 512, K = 2048 (gate/up) 3.90 -> 3.34 us (1.11x) KSplit 16, ColTiles 1 +// N = 2048, K = 512 (down) 4.16 -> 2.97 us (1.40x) KSplit 4, ColTiles 1 +Fp4MmaConfig PickFp4MmaConfig(int n, int k) { + static int const sm_count = onnxruntime::llm::common::getMultiProcessorCount(); + const int windows = k >> 7; // >= 1; the launcher only takes this path when k % 128 == 0 + const int col_tiles = (n + 15) / 16; + + if ((col_tiles + 3) / 4 >= sm_count) { + return {std::min(2, windows), 4}; + } + // Column-starved: give every block as many K-split warps as there are windows, up to the + // 16-warp (512-thread) block ceiling. + int k_split = 1; + while (k_split < 16 && k_split * 2 <= windows) { + k_split <<= 1; + } + return {k_split, 1}; +} + } // namespace #endif // CUDA_VERSION >= 12080 @@ -464,6 +747,7 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, int k, int block_size, bool is_bf16, + int sm_major, cudaStream_t stream) { #if defined(CUDA_VERSION) && CUDA_VERSION >= 12080 if (m <= 0 || n <= 0 || k <= 0) { @@ -475,6 +759,63 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, ORT_RETURN_IF_NOT(block_size == 16, "MatMulBlockQuantizedFp4Weight GEMV requires block_size == 16, got ", block_size, "."); ORT_RETURN_IF_NOT(k % 32 == 0, "MatMulBlockQuantizedFp4Weight GEMV requires K divisible by 32, got ", k, "."); const int k_blocks = (k + block_size - 1) / block_size; + + // Tensor-core sub-path: needs mma.m16n8k16 (SM80+), a whole number of 128-element K windows, + // and M within the mma's 8-row N extent. Everything else falls back to the scalar GEMV. + if (sm_major >= 8 && k % 128 == 0 && m <= 8 && Fp4GemvMmaEnabled()) { + const Fp4MmaConfig cfg = PickFp4MmaConfig(n, k); + const int cols_per_block = 16 * cfg.col_tiles; + const dim3 mma_threads{32, static_cast(cfg.k_split * cfg.col_tiles)}; + const dim3 mma_blocks{static_cast((n + cols_per_block - 1) / cols_per_block)}; + const uint8_t* mbp = reinterpret_cast(b_packed); + const uint8_t* mws = reinterpret_cast(weight_scale); + +#define ORT_LAUNCH_FP4_MMA_GEMV(T, KS, CT) \ + MatMulBlockQuantizedFp4WeightMmaGemvKernel \ + <<>>(reinterpret_cast(y), \ + reinterpret_cast(a), mbp, mws, \ + weight_scale_2, reinterpret_cast(bias), \ + m, n, k, k_blocks) + +#define ORT_DISPATCH_FP4_MMA_GEMV(T) \ + do { \ + if (cfg.col_tiles == 4) { \ + if (cfg.k_split >= 2) { \ + ORT_LAUNCH_FP4_MMA_GEMV(T, 2, 4); \ + } else { \ + ORT_LAUNCH_FP4_MMA_GEMV(T, 1, 4); \ + } \ + } else { \ + switch (cfg.k_split) { \ + case 16: \ + ORT_LAUNCH_FP4_MMA_GEMV(T, 16, 1); \ + break; \ + case 8: \ + ORT_LAUNCH_FP4_MMA_GEMV(T, 8, 1); \ + break; \ + case 4: \ + ORT_LAUNCH_FP4_MMA_GEMV(T, 4, 1); \ + break; \ + case 2: \ + ORT_LAUNCH_FP4_MMA_GEMV(T, 2, 1); \ + break; \ + default: \ + ORT_LAUNCH_FP4_MMA_GEMV(T, 1, 1); \ + break; \ + } \ + } \ + } while (0) + + if (is_bf16) { + ORT_DISPATCH_FP4_MMA_GEMV(nv_bfloat16); + } else { + ORT_DISPATCH_FP4_MMA_GEMV(half); + } +#undef ORT_DISPATCH_FP4_MMA_GEMV +#undef ORT_LAUNCH_FP4_MMA_GEMV + return CUDA_CALL(cudaGetLastError()); + } + const int rows_per_block = Fp4GemvRowsPerBlock(m, n); const dim3 threads{32, kGemvWarpsPerBlock}; const dim3 blocks{static_cast((n + kGemvWarpsPerBlock - 1) / kGemvWarpsPerBlock), @@ -522,6 +863,7 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, ORT_UNUSED_PARAMETER(k); ORT_UNUSED_PARAMETER(block_size); ORT_UNUSED_PARAMETER(is_bf16); + ORT_UNUSED_PARAMETER(sm_major); ORT_UNUSED_PARAMETER(stream); return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MatMulBlockQuantizedFp4Weight requires CUDA 12.8 or newer for NVFP4 support."); #endif diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h index fc3ef63f5bd46..4f30ef7943b40 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h @@ -60,6 +60,8 @@ Status LaunchAddBiasNvFp4(void* y, // uint8 (raw E4M3 bytes), weight_scale_2 is a device fp32 scalar, bias is an optional [N] vector // (may be null). Output y is [M, N] in the activation type. Requires block_size == 16 and // k % 32 == 0. Runs on any architecture with NVFP4 conversion intrinsics (CUDA >= 12.8). +// sm_major selects the tensor-core sub-path (mma.m16n8k16, SM80+); see the kernel comment in +// matmul_block_scaled_fp4.cu. Lower architectures use the scalar warp-reduction path. Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, const void* a, const void* b_packed, @@ -71,6 +73,7 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, int k, int block_size, bool is_bf16, + int sm_major, cudaStream_t stream); // Repacks the [N, ceil(K/block_size)] row-major E4M3 weight-scale tensor into the swizzled layout diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc index 32fa2db3e1045..eda3c65864b87 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc @@ -448,6 +448,160 @@ TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvDecodeRowTiledBf16) { } } +// Exercises the tensor-core GEMV sub-path (mma.m16n8k16), which the launcher selects on SM80+ +// when K is a multiple of 128 and M <= 8. None of the tests above reach it (they use K = 32/64). +// +// That path permutes the K axis so every lane's loads are contiguous, and folds the per-block +// E4M3 scale into the decoded weight instead of flushing the accumulator per block. A K +// permutation applied consistently to both operands is a no-op, so to catch an *inconsistent* +// one the weight, the block scales and the activation must all vary along K -- a uniform +// weight (as in the row-tiled tests above) would pass even with a broken k mapping. +// +// The shapes cover: a ragged N that is not a multiple of the 16-column warp tile, single and +// multi window K (128 -> 1 window, 256 -> 2, 512 -> 4), and an N large enough that the launcher +// picks the wide ColTiles = 4 / KSplit = 2 shape rather than the column-starved KSplit ladder. +namespace { + +// FP4 codes cycling along K: +1.0, -2.0, +0.5, +1.5 (E2M1 codes 0x2, 0xC, 0x1, 0x3). Negative +// values exercise the sign path of the prmt-based decode. +constexpr uint8_t kMmaCodes[4] = {0x2, 0xC, 0x1, 0x3}; +constexpr float kMmaValues[4] = {1.0f, -2.0f, 0.5f, 1.5f}; + +// The column rotates the cycle so neighbouring output columns are not interchangeable. +int MmaCodeIndex(int64_t col, int64_t kk) { return static_cast((kk + col) & 3); } + +// Block scale alternates 1.0 (E4M3 0x38) and 2.0 (0x40) along both K and N. +float MmaScaleValue(int64_t col, int64_t blk) { return ((col + blk) % 2 == 0) ? 1.0f : 2.0f; } +uint8_t MmaScaleByte(int64_t col, int64_t blk) { return ((col + blk) % 2 == 0) ? 0x38u : 0x40u; } + +float MmaActValue(int64_t row, int64_t kk) { return static_cast((kk % 4) + 1 + row); } + +// Every product is a multiple of 0.5 and every partial sum stays far below 2^23, so the fp32 +// reference below is exact and independent of summation order -- the kernel's permuted, split-K +// accumulation must match it bit for bit before the final cast to the output type. +void MakeMmaCase(int64_t m, int64_t n, int64_t k, + std::vector& b, std::vector& weight_scale, + std::vector& a, std::vector& expected) { + const int64_t k_blocks = k / 16; + b.assign(n * (k / 2), 0); + weight_scale.assign(n * k_blocks, 0); + for (int64_t col = 0; col < n; ++col) { + for (int64_t kk = 0; kk < k; kk += 2) { + const uint8_t lo = kMmaCodes[MmaCodeIndex(col, kk)]; + const uint8_t hi = kMmaCodes[MmaCodeIndex(col, kk + 1)]; + b[col * (k / 2) + kk / 2] = static_cast(lo | (hi << 4)); + } + for (int64_t blk = 0; blk < k_blocks; ++blk) { + weight_scale[col * k_blocks + blk] = MmaScaleByte(col, blk); + } + } + + a.assign(m * k, 0.0f); + for (int64_t row = 0; row < m; ++row) { + for (int64_t kk = 0; kk < k; ++kk) { + a[row * k + kk] = MmaActValue(row, kk); + } + } + + expected.assign(m * n, 0.0f); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < n; ++col) { + float sum = 0.0f; + for (int64_t kk = 0; kk < k; ++kk) { + sum += kMmaValues[MmaCodeIndex(col, kk)] * MmaScaleValue(col, kk / 16) * + MmaActValue(row, kk); + } + expected[row * n + col] = sum; + } + } +} + +struct MmaShape { + int64_t n; + int64_t k; +}; + +// N = 8704 gives 544 column tiles, i.e. 136 blocks at ColTiles = 4, which is above the SM count +// of every current device and so selects the wide shape. +constexpr MmaShape kMmaShapes[] = {{40, 128}, {512, 256}, {2048, 512}, {8704, 256}}; + +} // namespace + +TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvTensorCoreTilesFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + for (const auto& shape : kMmaShapes) { + for (int m_val : {1, 3, 4, 8}) { + const int64_t m = m_val; + const int64_t n = shape.n; + const int64_t k = shape.k; + SCOPED_TRACE("N = " + std::to_string(n) + ", K = " + std::to_string(k) + + ", M = " + std::to_string(m)); + + std::vector b, weight_scale; + std::vector a, expected; + MakeMmaCase(m, n, k, b, weight_scale, a, expected); + + OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", 16); + test.AddInput("A", {m, k}, FloatsToMLFloat16s(a)); + test.AddInput("B", {n, k / 2}, b); + test.AddInput("weight_scale", {n, k / 16}, weight_scale); + test.AddInput("weight_scale_2", {1}, {1.0f}); + test.AddOutput("Y", {m, n}, FloatsToMLFloat16s(expected)); + test.SetOutputTolerance(0.5f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } + } +} + +TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvTensorCoreTilesBf16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t n = 512; + constexpr int64_t k = 256; + + for (int m_val : {1, 4, 8}) { + const int64_t m = m_val; + SCOPED_TRACE("M = " + std::to_string(m)); + + std::vector b, weight_scale; + std::vector a, expected; + MakeMmaCase(m, n, k, b, weight_scale, a, expected); + + // Per-column bias, folded in by the kernel's store. + std::vector bias(n); + for (int64_t col = 0; col < n; ++col) { + bias[col] = static_cast((col % 5) - 2); + for (int64_t row = 0; row < m; ++row) { + expected[row * n + col] += bias[col]; + } + } + + OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", 16); + test.AddInput("A", {m, k}, FloatsToBFloat16s(a)); + test.AddInput("B", {n, k / 2}, b); + test.AddInput("weight_scale", {n, k / 16}, weight_scale); + test.AddInput("weight_scale_2", {1}, {1.0f}); + test.AddOptionalInputEdge(); // input_scale (skipped) + test.AddInput("bias", {n}, FloatsToBFloat16s(bias)); + test.AddOutput("Y", {m, n}, FloatsToBFloat16s(expected)); + test.SetOutputTolerance(0.5f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + } +} + #endif // USE_CUDA && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 } // namespace onnxruntime::test From 9d4e7ae63d18c2067f772d3ff636db8771e86a91 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 30 Jul 2026 20:59:08 +0000 Subject: [PATCH 6/9] lintrunner --- .../cuda/math/matmul_block_scaled_fp4.cu | 59 +++++++++---------- 1 file changed, 29 insertions(+), 30 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu index 10b0c80875bdd..f6c57c9e6b03e 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -229,17 +229,16 @@ struct Fp4Cvt { }; template -__global__ __launch_bounds__(32 * kGemvWarpsPerBlock, GemvMinBlocksPerSm::value) - void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, - const T* __restrict__ a, - const uint8_t* __restrict__ b_packed, - const uint8_t* __restrict__ weight_scale, - const float* __restrict__ weight_scale_2, - const T* __restrict__ bias, - int m, - int n, - int k, - int k_blocks) { +__global__ __launch_bounds__(32 * kGemvWarpsPerBlock, GemvMinBlocksPerSm::value) void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, + const T* __restrict__ a, + const uint8_t* __restrict__ b_packed, + const uint8_t* __restrict__ weight_scale, + const float* __restrict__ weight_scale_2, + const T* __restrict__ bias, + int m, + int n, + int k, + int k_blocks) { using Cvt = Fp4Cvt; using T2 = typename Cvt::T2; @@ -823,25 +822,25 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, const uint8_t* bp = reinterpret_cast(b_packed); const uint8_t* ws = reinterpret_cast(weight_scale); -#define ORT_DISPATCH_FP4_GEMV(T) \ - do { \ - switch (rows_per_block) { \ - case 4: \ - MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ - reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ - reinterpret_cast(bias), m, n, k, k_blocks); \ - break; \ - case 2: \ - MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ - reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ - reinterpret_cast(bias), m, n, k, k_blocks); \ - break; \ - default: \ - MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ - reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ - reinterpret_cast(bias), m, n, k, k_blocks); \ - break; \ - } \ +#define ORT_DISPATCH_FP4_GEMV(T) \ + do { \ + switch (rows_per_block) { \ + case 4: \ + MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ + reinterpret_cast(bias), m, n, k, k_blocks); \ + break; \ + case 2: \ + MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ + reinterpret_cast(bias), m, n, k, k_blocks); \ + break; \ + default: \ + MatMulBlockQuantizedFp4WeightGemvKernel<<>>( \ + reinterpret_cast(y), reinterpret_cast(a), bp, ws, weight_scale_2, \ + reinterpret_cast(bias), m, n, k, k_blocks); \ + break; \ + } \ } while (0) if (is_bf16) { From 0d480e6578a6e529e1dc7ac5d6aad1a2e0af2f34 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 31 Jul 2026 14:53:00 -0700 Subject: [PATCH 7/9] test(cuda): probe mma lane ownership in block-scaled GEMV Address PR review on the FP4/FP8 tensor-core decode GEMV: document the exact per-lane m16n8k16 fragment ownership (A rows -> output columns, B column -> activation row, D -> output rows 2t/2t+1) in both kernels and both operator docs, and add a one-hot lane-ownership probe test per data type so any swap of the row/column mapping fails immediately instead of cancelling out. Also make the FP8 tensor-core launch gate check M <= 8 explicitly, matching the FP4 gate, so the structural bound is enforced where it is documented. --- .../cuda/matmul_block_scaled_fp4.md | 7 ++ .../cuda/matmul_block_scaled_fp8.md | 21 ++++-- .../cuda/math/matmul_block_scaled_fp4.cu | 30 +++++++- .../cuda/math/matmul_block_scaled_fp8.cu | 38 ++++++++-- .../matmul_block_scaled_fp4_test.cc | 75 +++++++++++++++++++ .../matmul_block_scaled_fp8_test.cc | 63 ++++++++++++++++ 6 files changed, 220 insertions(+), 14 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md index c629131c994a4..3c834869246de 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md @@ -154,6 +154,13 @@ The fragments are free because the **weight goes in the mma `A` slot** and the `B`-col-major fragment. No transpose, no `ldmatrix`, no shared-memory staging. The mma `M` extent becomes the column count (16) and the mma `N` extent becomes `M`. +A lane does not own a whole dot product: with `g = lane >> 2` and `t = lane & 3`, +it supplies the weights of output columns `g` and `g + 8` (the `A` fragment) and +activation row `g` (the `B` fragment), and the accumulator it receives back covers +output rows `2t` and `2t + 1` of those two columns. Loads are therefore keyed off +`g` and stores off `t`; the kernel source carries the full fragment table, and the +`GemvTensorCoreLaneOwnership*` tests probe the mapping with a one-hot activation. + The K axis is then permuted so the four k-slots a lane needs are contiguous in memory, which is legal because K is a reduction axis and the same permutation is applied to both operands. A window is 128 K elements = 64 packed bytes; lane diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md index b83ed13bcc6fa..d9a5c292082d5 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md @@ -135,6 +135,7 @@ When the device is SM80 or newer and, in addition to the predicate above, - `K % 64 == 0` and `K >= 256`, - `block_size % 64 == 0`, +- `M <= 8` (the mma "N" extent), the GEMV runs `MatMulBlockScaledFp8MmaGemvKernel`, which replaces the FP32 FMA dot products with `mma.m16n8k16` (FP32 accumulate). The FMA kernel is ALU bound @@ -144,13 +145,19 @@ the tensor-core kernel cuts instructions per weight byte roughly 10x. The weight is fed to the mma **A** operand and the activation to the **B** operand, so both fragments match the tensors' natural row-major layouts. The mma "M" extent is therefore the output column count (16 columns per warp) and the mma -"N" extent is `M` (up to 8 rows). Fragment loads are made fully coalesced by -permuting the K axis - K is a reduction axis, so any permutation applied to both -operands leaves the result unchanged - which lets each lane load one contiguous -`uint4` of weight bytes per 64-element K window and feed four mma instructions -with it. `KSplit` warps per block take a strided share of the K windows and are -reduced through shared memory, which restores the memory-level parallelism lost -by giving each warp 16 columns instead of 1-4. +"N" extent is `M` (up to 8 rows). A lane does not own a whole dot product: with +`g = lane >> 2` and `t = lane & 3`, it supplies the weights of output columns +`g` and `g + 8` (the A fragment) and activation row `g` (the B fragment), and the +accumulator it receives back covers output rows `2t` and `2t + 1` of those two +columns. Loads are therefore keyed off `g` and stores off `t`; the kernel source +carries the full fragment table, and the `GemvTensorCoreLaneOwnership*` tests +probe the mapping with a one-hot activation. Fragment loads are made fully +coalesced by permuting the K axis - K is a reduction axis, so any permutation +applied to both operands leaves the result unchanged - which lets each lane load +one contiguous `uint4` of weight bytes per 64-element K window and feed four mma +instructions with it. `KSplit` warps per block take a strided share of the K +windows and are reduced through shared memory, which restores the memory-level +parallelism lost by giving each warp 16 columns instead of 1-4. Accuracy is unchanged: E4M3 to FP16/BF16 is lossless, the products are exact and the mma accumulates in FP32 just like the FMA path. Set `ORT_FP8_GEMV_MMA=0` to diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu index f6c57c9e6b03e..ef7e59d95cf7d 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -358,6 +358,25 @@ __global__ __launch_bounds__(32 * kGemvWarpsPerBlock, GemvMinBlocksPerSm> 2, t = lane & 3). A lane does NOT +// compute a whole dot product on its own: the tensor core exchanges partial products across the +// warp, so the A rows, the B column and the D rows a lane holds are three independent slices. +// +// ra[0] = (mma row g, mma k 2t, 2t+1) <- weight column col_lo +// ra[1] = (mma row g + 8, mma k 2t, 2t+1) <- weight column col_hi = col_lo + 8 +// ra[2] = (mma row g, mma k 2t+8, 2t+9) <- weight column col_lo +// ra[3] = (mma row g + 8, mma k 2t+8, 2t+9) <- weight column col_hi +// rb[0] = (mma k 2t, 2t+1, mma col g) <- activation row g +// rb[1] = (mma k 2t+8, 2t+9, mma col g) <- activation row g +// acc[0], acc[1] = (mma row g, mma col 2t, 2t+1) -> y[2t][col_lo], y[2t+1][col_lo] +// acc[2], acc[3] = (mma row g + 8, mma col 2t, 2t+1) -> y[2t][col_hi], y[2t+1][col_hi] +// +// So `g` selects both the output-column pair (through the A fragment) and the activation row +// (through the B fragment), while the accumulator the lane ends up owning covers output rows 2t +// and 2t + 1. Loads keyed off `g` and stores keyed off `t` are therefore both correct and are +// deliberately different; see the GemvTensorCoreLaneOwnership* tests for a one-hot probe of this +// mapping. +// // The K axis is then permuted so that the four k-slots a lane needs per mma step are // contiguous in memory. This is legal because K is a reduction axis and the same // permutation is applied to both operands. A window is 128 K elements = 64 packed bytes; @@ -476,8 +495,11 @@ __global__ __launch_bounds__(32 * KSplit * ColTiles, 1) void MatMulBlockQuantize using T2 = typename Cvt::T2; const int lane = threadIdx.x; - const int g = lane >> 2; // mma A row -> output column within the 16-column tile - const int t = lane & 3; // mma k-slot -> which 32-element slice of the window + // See the fragment-ownership table above: `g` indexes the A rows (output columns col_lo/col_hi) + // and the B column (activation row); `t` indexes the 32-element K slice and, in the accumulator, + // the output rows 2t / 2t + 1 this lane stores. + const int g = lane >> 2; + const int t = lane & 3; const int warp = static_cast(threadIdx.y); const int warp_k = (KSplit == 1) ? 0 : (warp % KSplit); const int warp_c = (KSplit == 1) ? warp : (warp / KSplit); @@ -761,6 +783,10 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, // Tensor-core sub-path: needs mma.m16n8k16 (SM80+), a whole number of 128-element K windows, // and M within the mma's 8-row N extent. Everything else falls back to the scalar GEMV. + // + // The M <= 8 bound is structural: one warp owns a 16-column x 8-row output tile, where lane + // (g = lane >> 2, t = lane & 3) reads activation row g, covers weight columns 16 * tile + g and + // + 8, and stores output rows 2t and 2t + 1 of those two columns. if (sm_major >= 8 && k % 128 == 0 && m <= 8 && Fp4GemvMmaEnabled()) { const Fp4MmaConfig cfg = PickFp4MmaConfig(n, k); const int cols_per_block = 16 * cfg.col_tiles; 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 1c81f67f9b49e..2bbfae44e5948 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu @@ -393,8 +393,28 @@ __global__ void MatMulBlockScaledFp8GemvKernel(AType* __restrict__ output, // mma D[16, 8] -> y[16 output columns][8 rows] // // The mma "M" extent is therefore our N (16 output columns per warp) and the mma "N" extent is -// our M (up to 8 rows). At M == 4 half the mma N lanes are idle, but the instruction count per -// weight byte still drops about 10x, which is what the ALU-bound path needs. +// our M (up to 8 rows, which is why the caller caps the GEMV at M <= 8). At M == 4 half the mma N +// lanes are idle, but the instruction count per weight byte still drops about 10x, which is what +// the ALU-bound path needs. +// +// Per-lane fragment ownership (PTX m16n8k16 layout, g = lane >> 2, t = lane & 3). A lane does NOT +// compute a whole dot product on its own: the tensor core exchanges partial products across the +// warp, so the A rows, the B column and the D rows a lane holds are three independent slices. +// +// ra[0] = (mma row g, mma k 2t, 2t+1) <- weight column col_lo = 16 * blockIdx.x + g +// ra[1] = (mma row g + 8, mma k 2t, 2t+1) <- weight column col_hi = col_lo + 8 +// ra[2] = (mma row g, mma k 2t+8, 2t+9) <- weight column col_lo +// ra[3] = (mma row g + 8, mma k 2t+8, 2t+9) <- weight column col_hi +// rb[0] = (mma k 2t, 2t+1, mma col g) <- activation row g +// rb[1] = (mma k 2t+8, 2t+9, mma col g) <- activation row g +// acc[0], acc[1] = (mma row g, mma col 2t, 2t+1) -> y[2t][col_lo], y[2t+1][col_lo] +// acc[2], acc[3] = (mma row g + 8, mma col 2t, 2t+1) -> y[2t][col_hi], y[2t+1][col_hi] +// +// So `g` selects both the output-column pair (through the A fragment) and the activation row +// (through the B fragment), while the accumulator the lane ends up owning covers output rows 2t +// and 2t + 1. Loads keyed off `g` and stores keyed off `t` are therefore both correct and are +// deliberately different; see the GemvTensorCoreLaneOwnership* tests for a one-hot probe of this +// mapping. // // Fragment loads are made fully coalesced by PERMUTING the K axis. K is a reduction axis, so any // permutation applied to BOTH operands leaves the result unchanged. Within a 64-element K window, @@ -486,8 +506,11 @@ __global__ void MatMulBlockScaledFp8MmaGemvKernel(AType* __restrict__ output, const int lane = threadIdx.x; const int warp = threadIdx.y; - const int g = lane >> 2; // mma group: output-column offset in the tile AND activation row - const int t = lane & 3; // mma thread-in-group: k sub-offset + // See the fragment-ownership table above: `g` indexes the A rows (output columns col_lo/col_hi) + // and the B column (activation row); `t` indexes the k sub-slice and, in the accumulator, the + // output rows 2t / 2t + 1 this lane stores. + const int g = lane >> 2; + const int t = lane & 3; const int windows = k >> 6; const int col_lo = blockIdx.x * 16 + g; @@ -766,7 +789,12 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, // Tensor-core path (SM80+). Beats the FMA kernel at every M on H200: 1.06-1.23x at M == 1 and // 1.4-1.87x at M == 4, where the FMA kernel is ALU bound. Needs 64-element K windows, and at // least 4 of them so KSplit warps have something to do. - if (sm_major >= 8 && k % 64 == 0 && k >= 256 && block_size % 64 == 0 && Fp8GemvMmaEnabled()) { + // + // M <= 8 is a hard requirement, not a heuristic: one warp owns a 16-column x 8-row output tile, + // where lane (g = lane >> 2, t = lane & 3) reads activation row g, covers weight columns + // 16 * blockIdx.x + g and + 8, and stores output rows 2t and 2t + 1 of those two columns. Rows + // beyond the mma's 8-row N extent have nowhere to live. + if (sm_major >= 8 && m <= 8 && 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) { diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc index eda3c65864b87..541094020838d 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc @@ -602,6 +602,81 @@ TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvTensorCoreTilesBf16) { } } +// 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 +// principle still land on a plausible value. Here the activation is one-hot, which collapses +// Y[row, col] to a single decoded weight times its block scale, and the probe offset differs per +// row. Both the weight code and the block scale cycle modulo 3 along N, so the two columns a lane +// owns (g and g + 8, which differ by 8 == 2 mod 3) never carry the same value, and every one of +// the 8 rows produces a different value in a given column. Any swap of the mma ownership -- rows +// 2t / 2t + 1 against columns g / g + 8, or the activation row against the output column -- shows +// up as a mismatch instead of cancelling out. +namespace { + +// E2M1 codes +1.0, -2.0, +0.5 and their values. +constexpr uint8_t kProbeCodes[3] = {0x2, 0xC, 0x1}; +constexpr float kProbeValues[3] = {1.0f, -2.0f, 0.5f}; +// E4M3 scale bytes 1.0, 4.0, 16.0 and their values. The 3 x 3 weight-times-scale products are all +// distinct, so no two (row, column) mappings can produce the same output value. +constexpr uint8_t kProbeScaleBytes[3] = {0x38, 0x48, 0x58}; +constexpr float kProbeScaleValues[3] = {1.0f, 4.0f, 16.0f}; + +int ProbeIndex(int64_t col, int64_t x) { return static_cast((col + x) % 3); } + +} // namespace + +TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvTensorCoreLaneOwnershipFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t m = 8; // full mma N extent: output rows 2t and 2t + 1 for every t + constexpr int64_t n = 40; // ragged: the last 16-column tile has only its low half in range + constexpr int64_t k = 128; // one K window + constexpr int64_t k_blocks = k / 16; + + std::vector b(n * (k / 2), 0); + std::vector weight_scale(n * k_blocks, 0); + for (int64_t col = 0; col < n; ++col) { + for (int64_t kk = 0; kk < k; kk += 2) { + const uint8_t lo = kProbeCodes[ProbeIndex(col, kk)]; + const uint8_t hi = kProbeCodes[ProbeIndex(col, kk + 1)]; + b[col * (k / 2) + kk / 2] = static_cast(lo | (hi << 4)); + } + for (int64_t blk = 0; blk < k_blocks; ++blk) { + weight_scale[col * k_blocks + blk] = kProbeScaleBytes[ProbeIndex(col, blk)]; + } + } + + // Row `row` probes K offset 11 * row. The eight (code index, scale-block index) pairs that + // produces are all different, and each row lands in a different 32-element lane slice of the + // window, so no two rows can be confused for one another. + std::vector a(m * k, 0.0f); + std::vector expected(m * n, 0.0f); + for (int64_t row = 0; row < m; ++row) { + const int64_t kk = 11 * row; + a[row * k + kk] = 1.0f; + for (int64_t col = 0; col < n; ++col) { + expected[row * n + col] = + kProbeValues[ProbeIndex(col, kk)] * kProbeScaleValues[ProbeIndex(col, kk / 16)]; + } + } + + OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", 16); + test.AddInput("A", {m, k}, FloatsToMLFloat16s(a)); + test.AddInput("B", {n, k / 2}, b); + test.AddInput("weight_scale", {n, k_blocks}, weight_scale); + test.AddInput("weight_scale_2", {1}, {1.0f}); + test.AddOutput("Y", {m, n}, FloatsToMLFloat16s(expected)); + test.SetOutputTolerance(0.01f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + #endif // USE_CUDA && defined(CUDA_VERSION) && CUDA_VERSION >= 12080 } // namespace onnxruntime::test 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 fb0842985bd3e..6f707e982dcd9 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -437,6 +437,69 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesBf16) { 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 +// principle still land on a plausible value. Here the activation is one-hot, which collapses +// Y[row, col] to a single weight times its block scale, and the probe offset differs per row. +// Both the weight and the block scale cycle modulo 3 along N, so the two columns a lane owns +// (g and g + 8, which differ by 8 == 2 mod 3) never carry the same value, and the 3 x 3 possible +// weight-times-scale products are all distinct, so the eight rows give eight different values in +// any given column. Any swap of the mma ownership -- output rows 2t / 2t + 1 against columns +// g / g + 8, or the activation row against the output column -- shows up as a mismatch instead of +// cancelling out. +TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreLaneOwnershipFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight."; + } + + constexpr int64_t m = 8; // full mma N extent: output rows 2t and 2t + 1 for every t + constexpr int64_t n = 40; // ragged: the last 16-column tile has only its low half in range + constexpr int64_t k = 256; + constexpr int64_t block_size = 64; + constexpr int64_t k_blocks = k / block_size; + + static const float kProbeWeights[3] = {1.0f, -2.0f, 0.5f}; // exact in E4M3 + static const float kProbeScales[3] = {1.0f, 4.0f, 16.0f}; // all 9 products are distinct + + 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(kProbeWeights[(col + i) % 3]); + } + for (int64_t kb = 0; kb < k_blocks; ++kb) { + b_scale[static_cast(col * k_blocks + kb)] = kProbeScales[(col + kb) % 3]; + } + } + + // Row `row` probes K offset 23 * row. The eight (weight index, scale-block index) pairs that + // produces are all different, and each row lands in a different 16-element lane slice of its + // 64-element K window, so no two rows can be confused for one another. + std::vector a(static_cast(m * k), 0.0f); + std::vector expected(static_cast(m * n), 0.0f); + for (int64_t row = 0; row < m; ++row) { + const int64_t i = 23 * row; + a[static_cast(row * k + i)] = 1.0f; + for (int64_t col = 0; col < n; ++col) { + expected[static_cast(row * n + col)] = + kProbeWeights[(col + i) % 3] * kProbeScales[(col + i / block_size) % 3]; + } + } + + 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.01f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultCudaExecutionProvider()); + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + #endif // USE_CUDA && !DISABLE_FLOAT8_TYPES && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 } // namespace onnxruntime::test From 1ea28c0a88dce392824b3bd1f082d0704f52b272 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 31 Jul 2026 23:54:05 +0000 Subject: [PATCH 8/9] fix(cuda): address review feedback on block-scaled decode GEMV - Parse the GEMV kill switches as bool so ORT_FP4_GEMV_MMA=false and friends actually disable the path instead of silently falling back to the default. - Pass the device properties into the GEMV launchers instead of caching the multiprocessor count in a function-static, which fixed the tiling heuristic to whichever GPU ran first in a heterogeneous multi-GPU process. This also unifies the sm_major derivation between the FP4 and FP8 call sites. - Cap __launch_bounds__ minBlocksPerSM at 4 for pre-SM80, where 1024 resident threads per SM makes 5/6 blocks of 256 unsatisfiable. - Raise the FP8 tensor-core test scales so |expected| is far above the output tolerance; the previous 1/256 scales left an all-zeros output passing. - Cover the FP8 RowsPerWarp = 8 FMA tile (M in [5, 8]) and the FP4 mma lo_ok == false predication (N = 36). - Align the docs: the FP8 mma path is not bit-identical to the FMA path, the M > 1 FMA dispatch re-tune changes the summation order, and FP4 row tiling is superseded by the mma path on SM80+ when K % 128 == 0. --- .../cuda/matmul_block_scaled_fp4.md | 4 ++ .../cuda/matmul_block_scaled_fp8.md | 8 ++- .../matmul_block_scaled_fp8_experiments.md | 36 +++++++++++ .../cuda/math/matmul_block_scaled_fp4.cc | 2 +- .../cuda/math/matmul_block_scaled_fp4.cu | 59 +++++++++++-------- .../cuda/math/matmul_block_scaled_fp4.h | 5 +- .../cuda/math/matmul_block_scaled_fp8.cc | 2 +- .../cuda/math/matmul_block_scaled_fp8.cu | 8 +-- .../cuda/math/matmul_block_scaled_fp8.h | 4 +- .../matmul_block_scaled_fp4_test.cc | 6 +- .../matmul_block_scaled_fp8_test.cc | 20 ++++--- 11 files changed, 108 insertions(+), 46 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md index 3c834869246de..a993ec5c5b866 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md @@ -137,6 +137,10 @@ Per-row fp32 accumulation order does not depend on `RowsPerBlock`, so results ar bit-identical across tilings. Set `ORT_FP4_GEMV_ROW_TILING=0` to force `RowsPerBlock == 1`. +Note that the tensor-core sub-path below takes precedence on SM80+ whenever +`K % 128 == 0`, which covers most production shapes. Row tiling is therefore what +actually runs on pre-SM80 devices or when `K` is not a multiple of 128. + ### Tensor-core sub-path (SM80+) On SM80 and newer, when `K % 128 == 0` and `M <= 8`, the warp reduction above is diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md index d9a5c292082d5..888ca0c1ee403 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp8.md @@ -159,9 +159,11 @@ instructions with it. `KSplit` warps per block take a strided share of the K windows and are reduced through shared memory, which restores the memory-level parallelism lost by giving each warp 16 columns instead of 1-4. -Accuracy is unchanged: E4M3 to FP16/BF16 is lossless, the products are exact and -the mma accumulates in FP32 just like the FMA path. Set `ORT_FP8_GEMV_MMA=0` to -fall back to the FMA kernel. +Every individual product is exact (E4M3 to FP16/BF16 is lossless) and the mma +accumulates in FP32 just like the FMA path, but the summation order differs, so +the result is **not** bit-identical to the FMA kernel. Note also that the mma +path is preferred over the FMA kernel by default on SM80+, including at `M == 1`. +Set `ORT_FP8_GEMV_MMA=0` to fall back to the FMA kernel. --- diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md index 99b39053d635c..03e707ff88cc5 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp8_experiments.md @@ -235,6 +235,11 @@ Below `N = 4096` the wider tiles leave too few warps to fill the GPU, and for `M > 1` the extra live registers (accumulators plus pre-issued loads) cost more than the added parallelism returns. Both measured slower, hence the guards. +> Superseded for `M > 1` by section 6.4, which re-tunes `ColsPerWarp` / `Unroll` +> for the speculative-decode tiles. Because `Unroll` changes the K chunk each +> lane accumulates first, the `M > 1` dispatch there is not bit-identical to +> `` (last-ulp only; the accumulation is still FP32). + ### 5.4 Results (H200, `M = 1`, CUDA graph, us, includes 0.68 us node overhead) | Shape (N x K) | cuBLAS FP16 | GEMV before | GEMV after | vs before | vs cuBLAS | @@ -346,6 +351,37 @@ per step, `M = 4`), CUDA graphs on: No other kernel family moved. This optimization is kept. +### 6.4 FMA fallback re-tune for `M > 1` + +The FMA kernel still runs when the tensor-core preconditions do not hold (pre-SM80, +`K < 256`, `K % 64 != 0` or `block_size % 64 != 0`), so the `M > 1` tiles were +re-tuned there as well. Widening `A` to FP32 is now hoisted out of the column loop +(one widening per row instead of one per row/column pair), which makes `ColsPerWarp` +profitable at `M > 1` for a second reason beyond memory-level parallelism: + +| Condition | Config | +|---|---| +| `2 <= M <= 2, N >= 8192` | `<2, 4, 1>` | +| `2 <= M <= 2, N >= 2048` | `<2, 2, 1>` | +| `2 <= M <= 2` otherwise | `<2, 1, 2>` | +| `3 <= M <= 4, N >= 4096` | `<4, 4, 1>` | +| `3 <= M <= 4, N >= 2048` | `<4, 2, 2>` | +| `3 <= M <= 4` otherwise | `<4, 1, 2>` | +| `M > 4` | `<8, 1, 1>` | + +Measured on H200 (us, `M = 4`, versus the previous `` and cuBLAS FP16): + +| Shape (N x K) | cuBLAS | `<4, 1, 1>` | tuned | +|---|---|---|---| +| 8192 x 2048 | 10.9 | 13.7 | **9.7** (`<4, 4, 1>`) | +| 4096 x 2048 | 8.3 | 8.1 | **6.9** (`<4, 4, 1>`) | +| 2048 x 4096 | 8.3 | 9.0 | **7.9** (`<4, 2, 2>`) | +| 512 x 2048 | 7.3 | 5.0 | **4.6** (`<4, 1, 2>`) | + +The hoisting itself is bit-identical (the per-lane `fmaf` sequence is unchanged), +but a different `Unroll` changes which K chunk a lane accumulates first, so the +re-tuned dispatch is a last-ulp change relative to ``. + --- ## 7. Benchmark Commands diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc index 13b9b80dfd230..8f5eb42d40862 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cc @@ -196,7 +196,7 @@ Status MatMulBlockQuantizedFp4Weight::ComputeImpl(OpKernelContext* context) cons k_i, SafeInt(block_size_), std::is_same::value, - sm_ / 10, + GetDeviceProp(), Stream(context)); } diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu index ef7e59d95cf7d..94fb1742737ee 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -14,7 +14,6 @@ #include #endif -#include "contrib_ops/cuda/llm/common/cuda_runtime_utils.h" #include "core/platform/env_var_utils.h" #include "core/providers/cuda/cuda_common.h" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" @@ -82,13 +81,22 @@ constexpr int kGemvWarpsPerBlock = 8; // fragments starts to cost more occupancy than the extra weight reuse buys. constexpr int kGemvMaxRowsPerBlock = 4; +// SM75 allows only 1024 resident threads per multiprocessor, i.e. 4 blocks of 256. Asking for +// more there is unsatisfiable and only makes ptxas over-restrict registers and spill. +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800 +constexpr int kGemvMaxBlocksPerSm = 4; +#else +constexpr int kGemvMaxBlocksPerSm = 8; +#endif + // Occupancy target for the GEMV. Holding RowsPerBlock accumulators costs registers, and left // unconstrained nvcc trades occupancy for scheduling freedom: the 4-row tile lands on 66 // registers, which drops the block from 4 to 3 per SM and costs ~7% on the large shapes. Pin a // register budget per tile instead; all three instantiations compile spill-free at these targets. template struct GemvMinBlocksPerSm { - static constexpr int value = (RowsPerBlock == 1) ? 6 : ((RowsPerBlock == 2) ? 5 : 4); + static constexpr int target = (RowsPerBlock == 1) ? 6 : ((RowsPerBlock == 2) ? 5 : 4); + static constexpr int value = target < kGemvMaxBlocksPerSm ? target : kGemvMaxBlocksPerSm; }; // ----------------------------------------------------------------------------- @@ -109,7 +117,12 @@ struct GemvMinBlocksPerSm { // the launcher only enables it when gridDim.x alone already fills the device; see // Fp4GemvRowsPerBlock() below. The per-row fp32 accumulation order is independent of // RowsPerBlock, so results are bit-identical across tilings. +// +// Note that on SM80+ the tensor-core kernel further below takes precedence whenever +// K % 128 == 0, so row tiling is what actually runs on pre-SM80 devices or when K is not a +// multiple of 128. // ----------------------------------------------------------------------------- + // ----------------------------------------------------------------------------- // Fast NVFP4 (E2M1) -> half / bfloat16 conversion. // @@ -229,16 +242,18 @@ struct Fp4Cvt { }; template -__global__ __launch_bounds__(32 * kGemvWarpsPerBlock, GemvMinBlocksPerSm::value) void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, - const T* __restrict__ a, - const uint8_t* __restrict__ b_packed, - const uint8_t* __restrict__ weight_scale, - const float* __restrict__ weight_scale_2, - const T* __restrict__ bias, - int m, - int n, - int k, - int k_blocks) { +__global__ __launch_bounds__(32 * kGemvWarpsPerBlock, + GemvMinBlocksPerSm::value) void MatMulBlockQuantizedFp4WeightGemvKernel( + T* __restrict__ y, + const T* __restrict__ a, + const uint8_t* __restrict__ b_packed, + const uint8_t* __restrict__ weight_scale, + const float* __restrict__ weight_scale_2, + const T* __restrict__ bias, + int m, + int n, + int k, + int k_blocks) { using Cvt = Fp4Cvt; using T2 = typename Cvt::T2; @@ -608,7 +623,7 @@ __global__ __launch_bounds__(32 * KSplit * ColTiles, 1) void MatMulBlockQuantize bool Fp4GemvRowTilingEnabled() { static bool const enabled = - onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_ROW_TILING", 1) == 1; + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_ROW_TILING", true); return enabled; } @@ -625,11 +640,10 @@ bool Fp4GemvRowTilingEnabled() { // N = 512 (shared gate) 64 col blocks 3.47 -> 4.27 us (0.81x, fewer blocks than SMs) // // So gate on the column grid covering at least one full wave of SMs on its own. -int Fp4GemvRowsPerBlock(int m, int n) { +int Fp4GemvRowsPerBlock(int m, int n, int sm_count) { if (m <= 1 || !Fp4GemvRowTilingEnabled()) { return 1; } - static int const sm_count = onnxruntime::llm::common::getMultiProcessorCount(); const int col_blocks = (n + kGemvWarpsPerBlock - 1) / kGemvWarpsPerBlock; if (col_blocks < sm_count) { return 1; @@ -641,7 +655,7 @@ int Fp4GemvRowsPerBlock(int m, int n) { bool Fp4GemvMmaEnabled() { static bool const enabled = - onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_MMA", 1) == 1; + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP4_GEMV_MMA", true); return enabled; } @@ -661,8 +675,7 @@ struct Fp4MmaConfig { // N = 248320, K = 2048 (lm_head) 538.5 -> 107.7 us (5.00x) KSplit 2, ColTiles 4 // N = 512, K = 2048 (gate/up) 3.90 -> 3.34 us (1.11x) KSplit 16, ColTiles 1 // N = 2048, K = 512 (down) 4.16 -> 2.97 us (1.40x) KSplit 4, ColTiles 1 -Fp4MmaConfig PickFp4MmaConfig(int n, int k) { - static int const sm_count = onnxruntime::llm::common::getMultiProcessorCount(); +Fp4MmaConfig PickFp4MmaConfig(int n, int k, int sm_count) { const int windows = k >> 7; // >= 1; the launcher only takes this path when k % 128 == 0 const int col_tiles = (n + 15) / 16; @@ -768,7 +781,7 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, int k, int block_size, bool is_bf16, - int sm_major, + const cudaDeviceProp& device_prop, cudaStream_t stream) { #if defined(CUDA_VERSION) && CUDA_VERSION >= 12080 if (m <= 0 || n <= 0 || k <= 0) { @@ -787,8 +800,8 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, // The M <= 8 bound is structural: one warp owns a 16-column x 8-row output tile, where lane // (g = lane >> 2, t = lane & 3) reads activation row g, covers weight columns 16 * tile + g and // + 8, and stores output rows 2t and 2t + 1 of those two columns. - if (sm_major >= 8 && k % 128 == 0 && m <= 8 && Fp4GemvMmaEnabled()) { - const Fp4MmaConfig cfg = PickFp4MmaConfig(n, k); + if (device_prop.major >= 8 && k % 128 == 0 && m <= 8 && Fp4GemvMmaEnabled()) { + const Fp4MmaConfig cfg = PickFp4MmaConfig(n, k, device_prop.multiProcessorCount); const int cols_per_block = 16 * cfg.col_tiles; const dim3 mma_threads{32, static_cast(cfg.k_split * cfg.col_tiles)}; const dim3 mma_blocks{static_cast((n + cols_per_block - 1) / cols_per_block)}; @@ -841,7 +854,7 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, return CUDA_CALL(cudaGetLastError()); } - const int rows_per_block = Fp4GemvRowsPerBlock(m, n); + const int rows_per_block = Fp4GemvRowsPerBlock(m, n, device_prop.multiProcessorCount); const dim3 threads{32, kGemvWarpsPerBlock}; const dim3 blocks{static_cast((n + kGemvWarpsPerBlock - 1) / kGemvWarpsPerBlock), static_cast((m + rows_per_block - 1) / rows_per_block)}; @@ -888,7 +901,7 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, ORT_UNUSED_PARAMETER(k); ORT_UNUSED_PARAMETER(block_size); ORT_UNUSED_PARAMETER(is_bf16); - ORT_UNUSED_PARAMETER(sm_major); + ORT_UNUSED_PARAMETER(device_prop); ORT_UNUSED_PARAMETER(stream); return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MatMulBlockQuantizedFp4Weight requires CUDA 12.8 or newer for NVFP4 support."); #endif diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h index 4f30ef7943b40..190c1dfd3e7eb 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.h @@ -60,7 +60,8 @@ Status LaunchAddBiasNvFp4(void* y, // uint8 (raw E4M3 bytes), weight_scale_2 is a device fp32 scalar, bias is an optional [N] vector // (may be null). Output y is [M, N] in the activation type. Requires block_size == 16 and // k % 32 == 0. Runs on any architecture with NVFP4 conversion intrinsics (CUDA >= 12.8). -// sm_major selects the tensor-core sub-path (mma.m16n8k16, SM80+); see the kernel comment in +// device_prop selects the tensor-core sub-path (mma.m16n8k16, SM80+) and sizes the M-tiling and +// K-split heuristics from the multiprocessor count; see the kernel comment in // matmul_block_scaled_fp4.cu. Lower architectures use the scalar warp-reduction path. Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, const void* a, @@ -73,7 +74,7 @@ Status LaunchMatMulBlockQuantizedFp4WeightGemv(void* y, int k, int block_size, bool is_bf16, - int sm_major, + const cudaDeviceProp& device_prop, cudaStream_t stream); // Repacks the [N, ceil(K/block_size)] row-major E4M3 weight-scale tensor into the swizzled layout diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc index 172f534ba0ae8..018c7dfbd8da4 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc @@ -119,7 +119,7 @@ Status MatMulBlockQuantizedFp8Weight::ComputeImpl(OpKernelContext* context) cons k_i, SafeInt(block_size_), std::is_same::value, - GetDeviceProp().major, + GetDeviceProp(), Stream(context)); } 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 2bbfae44e5948..e0e7171bc7ddf 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cu @@ -630,7 +630,7 @@ __global__ void MatMulBlockScaledFp8MmaGemvKernel(AType* __restrict__ output, // 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", 1) == 1; + static bool const enabled = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_FP8_GEMV_MMA", true); return enabled; } @@ -771,7 +771,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, int k, int block_size, bool is_bf16, - int sm_major, + const cudaDeviceProp& device_prop, cudaStream_t stream) { #if !defined(DISABLE_FLOAT8_TYPES) && defined(CUDA_VERSION) && CUDA_VERSION >= 11080 if (m <= 0 || n <= 0 || k <= 0) { @@ -794,7 +794,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, // where lane (g = lane >> 2, t = lane & 3) reads activation row g, covers weight columns // 16 * blockIdx.x + g and + 8, and stores output rows 2t and 2t + 1 of those two columns. Rows // beyond the mma's 8-row N extent have nowhere to live. - if (sm_major >= 8 && m <= 8 && k % 64 == 0 && k >= 256 && block_size % 64 == 0 && Fp8GemvMmaEnabled()) { + if (device_prop.major >= 8 && m <= 8 && 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) { @@ -896,7 +896,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, ORT_UNUSED_PARAMETER(k); ORT_UNUSED_PARAMETER(block_size); ORT_UNUSED_PARAMETER(is_bf16); - ORT_UNUSED_PARAMETER(sm_major); + ORT_UNUSED_PARAMETER(device_prop); ORT_UNUSED_PARAMETER(stream); return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MatMulBlockQuantizedFp8Weight requires CUDA 11.8 or later."); #endif diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h index bba4c7ec9adc7..d4388520014c3 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.h @@ -67,7 +67,7 @@ Status LaunchQuantizeDequantizeActivationFp8(void* a_out, // directly (no [N, K] dequant buffer). a is [M, K] activation (FP16/BF16), b_fp8 is [N, K] // FP8 E4M3, weight_scale is [N, ceil(K/block_size)] fp32, bias is an optional [N] vector (may be // null). Output y is [M, N] in the activation type. Requires k % 16 == 0 and block_size % 16 == 0. -// sm_major selects the tensor-core (mma.m16n8k16) variant, which needs SM80+. +// device_prop selects the tensor-core (mma.m16n8k16) variant, which needs SM80+. // Runs on any architecture with FP8 conversion intrinsics (CUDA >= 11.8). Status LaunchMatMulBlockScaledFp8Gemv(void* y, const void* a, @@ -79,7 +79,7 @@ Status LaunchMatMulBlockScaledFp8Gemv(void* y, int k, int block_size, bool is_bf16, - int sm_major, + const cudaDeviceProp& device_prop, cudaStream_t stream); } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc index 541094020838d..22d0ddba08870 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc @@ -522,8 +522,10 @@ struct MmaShape { }; // N = 8704 gives 544 column tiles, i.e. 136 blocks at ColTiles = 4, which is above the SM count -// of every current device and so selects the wide shape. -constexpr MmaShape kMmaShapes[] = {{40, 128}, {512, 256}, {2048, 512}, {8704, 256}}; +// of every current device and so selects the wide shape. N = 36 leaves 4 columns in the last +// 16-column tile, which is the only way to make a lane's *low* column fall out of range (N = 40 +// only exercises the high column), so it covers the lo_ok == false predication. +constexpr MmaShape kMmaShapes[] = {{36, 128}, {40, 128}, {512, 256}, {2048, 512}, {8704, 256}}; } // namespace 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 6f707e982dcd9..fdcb6f84754a8 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc @@ -251,9 +251,10 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvSpeculativeDecodeTilesFp16) { GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp8Weight."; } - // m = 2 -> RowsPerWarp 2, m = 3/4 -> RowsPerWarp 4 (m = 3 also leaves a ragged row tail). - // n picks ColsPerWarp 1 (n < 2048), 2 and 4; all leave a ragged column tail. - for (const int64_t m : {2, 3, 4}) { + // m = 2 -> RowsPerWarp 2, m = 3/4 -> RowsPerWarp 4, m = 5/8 -> RowsPerWarp 8 (m = 3 and m = 5 + // also leave a ragged row tail). n picks ColsPerWarp 1 (n < 2048), 2 and 4; all leave a ragged + // column tail. + for (const int64_t m : {2, 3, 4, 5, 8}) { for (const int64_t n : {1026, 2050, 4098, 8194}) { constexpr int64_t k = 64; // K % 16 == 0 -> GEMV path; two 32-element K blocks constexpr int64_t block_size = 32; @@ -320,7 +321,10 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesFp16) { const int64_t k_blocks = c.k / c.block_size; for (const int64_t m : {1, 3, 4, 8}) { // Periods 3 (weight) and 4 (activation) are coprime, so no (row, col) pair sums to zero by - // symmetry. Scales are 1/256-based so the reference stays small enough to be exact in FP16. + // symmetry. Even so the signed terms cancel heavily, so the scales are kept in [0.25, 0.75] + // rather than scaled down: every product is a multiple of 1/8 and the reference stays exact + // in FP16, while |expected| stays well above the tolerance below (an all-zero output must + // not pass). std::vector b(static_cast(c.n * c.k)); std::vector b_ref(static_cast(c.n * c.k)); for (int64_t col = 0; col < c.n; ++col) { @@ -333,7 +337,7 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesFp16) { std::vector b_scale(static_cast(c.n * k_blocks)); for (int64_t col = 0; col < c.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) / 256.0f; + b_scale[static_cast(col * k_blocks + kb)] = static_cast(1 + (col + kb) % 3) / 4.0f; } } std::vector a(static_cast(m * c.k)); @@ -360,7 +364,7 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesFp16) { test.AddInput("B", {c.n, c.k}, b); test.AddInput("b_scale", {c.n, k_blocks}, b_scale); test.AddOutput("Y", {m, c.n}, FloatsToMLFloat16s(expected)); - test.SetOutputTolerance(0.05f); + test.SetOutputTolerance(0.005f); std::vector> execution_providers; execution_providers.push_back(DefaultCudaExecutionProvider()); @@ -397,7 +401,7 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesBf16) { 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) / 256.0f; + b_scale[static_cast(col * k_blocks + kb)] = static_cast(1 + (col + kb) % 3) / 4.0f; } } std::vector a(static_cast(m * k)); @@ -430,7 +434,7 @@ TEST(MatMulBlockQuantizedFp8WeightOpTest, GemvTensorCoreTilesBf16) { test.AddOptionalInputEdge(); // a_scale (skipped) test.AddInput("bias", {n}, FloatsToBFloat16s(bias)); test.AddOutput("Y", {m, n}, FloatsToBFloat16s(expected)); - test.SetOutputTolerance(0.1f); + test.SetOutputTolerance(0.02f); std::vector> execution_providers; execution_providers.push_back(DefaultCudaExecutionProvider()); From 0d1fd027c4ce1acb1f0f3be35b7f461d35824ed2 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 1 Aug 2026 00:40:47 +0000 Subject: [PATCH 9/9] lintrunner --- .../cuda/math/matmul_block_scaled_fp4.cu | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu index 94fb1742737ee..9736a6d5d15d4 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -243,17 +243,16 @@ struct Fp4Cvt { template __global__ __launch_bounds__(32 * kGemvWarpsPerBlock, - GemvMinBlocksPerSm::value) void MatMulBlockQuantizedFp4WeightGemvKernel( - T* __restrict__ y, - const T* __restrict__ a, - const uint8_t* __restrict__ b_packed, - const uint8_t* __restrict__ weight_scale, - const float* __restrict__ weight_scale_2, - const T* __restrict__ bias, - int m, - int n, - int k, - int k_blocks) { + GemvMinBlocksPerSm::value) void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, + const T* __restrict__ a, + const uint8_t* __restrict__ b_packed, + const uint8_t* __restrict__ weight_scale, + const float* __restrict__ weight_scale_2, + const T* __restrict__ bias, + int m, + int n, + int k, + int k_blocks) { using Cvt = Fp4Cvt; using T2 = typename Cvt::T2;