From ca8d3947c39e1b8f8f21379b68f7704716282744 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 17 Aug 2026 08:33:47 +0000 Subject: [PATCH 1/2] [CUDA] Vectorize the NVFP4 weight dequantization for prefill The prefill fallback of MatMulBlockQuantizedFp4Weight expands the packed weight into an [N, K] scratch buffer before calling cuBLAS. That expansion ran one thread per packed byte: a 1-byte load, two 2-byte stores, two integer divisions and a global load of weight_scale_2 per thread. DequantizeNvFp4Vec8Kernel gives each thread exactly one 8-element K chunk of one row, so a warp issues one contiguous 128-byte load and one contiguous 512-byte store. The row index comes from blockIdx.y instead of a 64-bit division, weight_scale_2 is hoisted into a register, and the scale index is advanced incrementally. Codes are decoded with the existing branch-free Fp4Cvt prmt lookup (added for the decode GEMV in #31155) rather than __nv_cvt_fp4x2_to_halfraw2(), which is emulated in software on SM90 with branches and a subnormal normalization loop. Fp4Cvt reproduces the intrinsic's bit pattern exactly, so the dequantized weight is bitwise identical to the scalar kernel. The scalar kernel is kept for odd block_size or K % 8 != 0. --- .../cuda/matmul_block_scaled_fp4.md | 31 +++ .../cuda/math/matmul_block_scaled_fp4.cu | 95 ++++++++- .../matmul_block_scaled_fp4_test.cc | 191 ++++++++++++++++++ 3 files changed, 315 insertions(+), 2 deletions(-) diff --git a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md index a993ec5c5b866..6c69b4d502475 100644 --- a/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md +++ b/docs/contrib_ops/cuda/matmul_block_scaled_fp4.md @@ -28,6 +28,7 @@ Source files: 3. [Dispatch Chain](#3-dispatch-chain) 4. [Decode Path - Fused GEMV](#4-decode-path---fused-gemv) 5. [Default Path - Dequantize + cuBLAS](#5-default-path---dequantize--cublas) + - [5.1 Vectorized dequantization](#51-vectorized-dequantization) 6. [Native SM120 FP4 x FP4 Path](#6-native-sm120-fp4-x-fp4-path) 7. [PrePack](#7-prepack) 8. [Environment Variables](#8-environment-variables) @@ -220,6 +221,31 @@ This path keeps full-precision activations and runs on CUDA devices with NVFP4 conversion intrinsic support in the configured CUDA toolkit. It is the default prefill path when the SM120 native environment variable is not enabled. +### 5.1 Vectorized dequantization + +When `K % 8 == 0` and `block_size` is even - the layout every real NVFP4 model +uses - `LaunchDequantizeNvFp4` picks `DequantizeNvFp4Vec8Kernel` instead of the +scalar kernel. Each thread owns exactly one 8-element K chunk of one row, so a +warp issues one contiguous 128-byte packed load and one contiguous 512-byte +store. Widening the per-thread chunk beyond one `uint4` store was measured to be +about 2x slower because each store instruction then strides across lanes +(1.9 vs 3.9 TB/s on H200). The row index comes from `blockIdx.y`, which removes +the 64-bit division of the scalar kernel, `weight_scale_2` is hoisted into a +register, and codes are decoded with the same branch-free `Fp4Cvt` `prmt` +lookup the decode GEMV uses rather than the software-emulated +`__nv_cvt_fp4x2_to_halfraw2()`. + +The output is bitwise identical to the scalar kernel. Measured on H200 for +`M = 1024`, BF16, `block_size = 16` (median dequant kernel time): + +| N | K | scalar | vectorized | speedup | +|---:|---:|---:|---:|---:| +| 4096 | 4096 | 60.7 us | 15.5 us | 3.93x | +| 6144 | 2048 | 46.1 us | 12.1 us | 3.81x | +| 2048 | 6144 | 46.2 us | 11.9 us | 3.88x | + +The scalar kernel remains for odd `block_size` or `K % 8 != 0`. + --- ## 6. Native SM120 FP4 x FP4 Path @@ -304,6 +330,11 @@ CUDA_VISIBLE_DEVICES=0 "$ORT_BUILD/onnxruntime_provider_test" \ --gtest_filter='MatMulBlockQuantizedFp4WeightOpTest.*' ``` +The `Gemv*` cases cover the decode path and the `PrefillDequant*` cases cover the +dequantize + cuBLAS path, with `M > 8` so the GEMV is skipped: `*Vectorized*` for +`DequantizeNvFp4Vec8Kernel` and `OddBlockSize` / `KNotMultipleOf8` for the scalar +fallback, each in both FP16 and BF16. + Python harness examples: ```bash 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 9736a6d5d15d4..2143f7820b527 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -241,6 +241,72 @@ struct Fp4Cvt { static __device__ __forceinline__ float2 ToFloat2(T2 v) { return Traits::to_float2(v); } }; +// Vectorized dequantization for K % 8 == 0 and even block_size. Each thread owns exactly one +// 8-element K chunk of a row, so a warp issues one contiguous 128-byte packed load and one +// contiguous 512-byte store; giving a thread a wider chunk instead would make every store +// instruction strided across lanes, which measures ~2x slower on H200 (1.9 vs 3.9 TB/s). +// Codes are decoded with the branch-free Fp4Cvt above rather than the software-emulated +// __nv_cvt_fp4x2_to_halfraw2(). The row index comes from a 2D grid (blockIdx.y), which removes the +// 64-bit idx / half_k division of the scalar kernel, and weight_scale_2 is hoisted into a register. +// Because block_size is even, both codes of a byte always share a scale, so the scale index only +// advances once per block_size/2 bytes - tracked incrementally instead of by per-element division. +// Fp4Cvt reproduces the intrinsic's bit pattern exactly, so the output is bitwise identical to +// DequantizeNvFp4Kernel. +template +__global__ void DequantizeNvFp4Vec8Kernel(T* __restrict__ out, + const uint8_t* __restrict__ b_packed, + const uint8_t* __restrict__ weight_scale, + const float* __restrict__ weight_scale_2, + int n, + int k, + int k_blocks, + int block_size, + int chunks_per_row) { + using Cvt = Fp4Cvt; + using T2 = typename Cvt::T2; + + const int chunk = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (chunk >= chunks_per_row) { + return; + } + // Scale bookkeeping is done in units of T2 pairs (= one packed byte). + const int pair0 = chunk << 2; // first packed byte of the chunk within the row + const int pairs_per_block = block_size >> 1; + const int blk0 = pair0 / pairs_per_block; + const int left0 = pairs_per_block - (pair0 - blk0 * pairs_per_block); + const float g = *weight_scale_2; + + for (int row = blockIdx.y; row < n; row += gridDim.y) { + const uint32_t word = reinterpret_cast(b_packed + static_cast(row) * (k >> 1))[chunk]; + const uint8_t* srow = weight_scale + static_cast(row) * k_blocks; + + T2 pairs[4]; + const uint32_t mag = word & 0x77777777u; + const uint32_t sgn = (word >> 3) & 0x11111111u; + Cvt::DecodeQuad(mag, sgn, pairs[0], pairs[1]); + Cvt::DecodeQuad(mag >> 16, sgn >> 16, pairs[2], pairs[3]); + + int blk = blk0; + int left = left0; + float s = e4m3_to_float(srow[blk]) * g; + + T outv[8]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const float2 v = Cvt::ToFloat2(pairs[j]); + outv[2 * j] = from_float(v.x * s); + outv[2 * j + 1] = from_float(v.y * s); + if (--left == 0) { + ++blk; + left = pairs_per_block; + s = (blk < k_blocks) ? e4m3_to_float(srow[blk]) * g : 0.f; + } + } + + reinterpret_cast(out + static_cast(row) * k)[chunk] = *reinterpret_cast(outv); + } +} + template __global__ __launch_bounds__(32 * kGemvWarpsPerBlock, GemvMinBlocksPerSm::value) void MatMulBlockQuantizedFp4WeightGemvKernel(T* __restrict__ y, @@ -710,11 +776,36 @@ Status LaunchDequantizeNvFp4(void* b_dequant, return Status::OK(); } const int k_blocks = (k + block_size - 1) / block_size; - constexpr int kThreads = 256; - const int blocks = static_cast((total + kThreads - 1) / kThreads); const uint8_t* bp = reinterpret_cast(b_packed); const uint8_t* ws = reinterpret_cast(weight_scale); + // Fast path: K is a multiple of 8 (so each thread owns an aligned 8-element chunk and both the + // packed load and the FP16/BF16 store are warp-contiguous) and block_size is even (so both codes + // of a packed byte share a scale). This is the common prefill layout. + if (k % 8 == 0 && block_size % 2 == 0) { + const int chunks_per_row = k >> 3; + int threads = 32; + while (threads < 256 && threads < chunks_per_row) { + threads <<= 1; + } + const unsigned int grid_x = static_cast((chunks_per_row + threads - 1) / threads); + const unsigned int grid_y = static_cast(n < 65535 ? n : 65535); + const dim3 blocks{grid_x, grid_y}; + if (is_bf16) { + DequantizeNvFp4Vec8Kernel<<>>( + reinterpret_cast(b_dequant), bp, ws, weight_scale_2, n, k, k_blocks, block_size, + chunks_per_row); + } else { + DequantizeNvFp4Vec8Kernel<<>>( + reinterpret_cast(b_dequant), bp, ws, weight_scale_2, n, k, k_blocks, block_size, + chunks_per_row); + } + return CUDA_CALL(cudaGetLastError()); + } + + constexpr int kThreads = 256; + const int blocks = static_cast((total + kThreads - 1) / kThreads); + if (is_bf16) { DequantizeNvFp4Kernel<<>>( reinterpret_cast(b_dequant), bp, ws, weight_scale_2, n, k, k_blocks, block_size); 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 22d0ddba08870..e366df6c34ade 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc @@ -679,6 +679,197 @@ TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvTensorCoreLaneOwnershipFp16) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +// --------------------------------------------------------------------------- +// Prefill (M > kGemvMaxM) dequantize + cuBLAS path. +// +// LaunchDequantizeNvFp4 picks DequantizeNvFp4Vec8Kernel when K % 8 == 0 and block_size is even, +// and the scalar DequantizeNvFp4Kernel otherwise. The cases below cover both sides of that guard; +// every one uses M > 8 so the decode GEMV is skipped and the dequant actually runs. +// --------------------------------------------------------------------------- + +// Vectorized path: K % 8 == 0, even block_size, several K blocks with different scales, and a +// negative weight row so the sign bit of the packed nibble is exercised. +TEST(MatMulBlockQuantizedFp4WeightOpTest, PrefillDequantVectorizedFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t m = 16; // > kGemvMaxM (8), so the dequant + cuBLAS path runs + constexpr int64_t n = 3; + constexpr int64_t k = 32; + constexpr int64_t k_blocks = k / 16; + + // Row 0 = +1.0 (0x22), row 1 = +2.0 (0x44), row 2 = -1.5 (sign 0x8 | 0x3 -> 0xBB). + std::vector b(n * (k / 2)); + for (int64_t j = 0; j < k / 2; ++j) { + b[0 * (k / 2) + j] = 0x22; + b[1 * (k / 2) + j] = 0x44; + b[2 * (k / 2) + j] = 0xBB; + } + // Row 0 blocks {1.0, 2.0}, row 1 {1.0, 1.0}, row 2 {2.0, 0.5}. + std::vector weight_scale = {0x38, 0x40, 0x38, 0x38, 0x40, 0x30}; + std::vector weight_scale_2 = {1.0f}; + + std::vector a(m * k); + for (int64_t row = 0; row < m; ++row) { + for (int64_t col = 0; col < k; ++col) { + a[row * k + col] = static_cast(row + 1); + } + } + // Each block contributes 16 * (row + 1) * value * block_scale: + // Y[r, 0] = 16 (r+1) (1*1 + 1*2) = 48 (r+1) + // Y[r, 1] = 16 (r+1) (2*1 + 2*1) = 64 (r+1) + // Y[r, 2] = 16 (r+1) (-1.5*2 + -1.5*0.5) = -60 (r+1) + std::vector expected(m * n); + for (int64_t row = 0; row < m; ++row) { + const float scale = static_cast(row + 1); + expected[row * n + 0] = 48.0f * scale; + expected[row * n + 1] = 64.0f * scale; + expected[row * n + 2] = -60.0f * scale; + } + + 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); +} + +// Vectorized path with BF16, a bias, a non-unit weight_scale_2, and M = 9 to pin the boundary +// just above the decode GEMV cutoff. +TEST(MatMulBlockQuantizedFp4WeightOpTest, PrefillDequantVectorizedBiasBf16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t m = 9; // kGemvMaxM + 1 + constexpr int64_t n = 2; + constexpr int64_t k = 64; + constexpr int64_t k_blocks = k / 16; + + // Row 0 = +1.0 (0x22), row 1 = -2.0 (sign 0x8 | 0x4 -> 0xCC). + std::vector b(n * (k / 2)); + for (int64_t j = 0; j < k / 2; ++j) { + b[0 * (k / 2) + j] = 0x22; + b[1 * (k / 2) + j] = 0xCC; + } + // Row 0 blocks {1.0, 2.0, 1.0, 2.0}, row 1 all 1.0. + std::vector weight_scale = {0x38, 0x40, 0x38, 0x40, 0x38, 0x38, 0x38, 0x38}; + std::vector weight_scale_2 = {0.5f}; + std::vector bias = {1.0f, 2.0f}; + + std::vector a(m * k, 1.0f); + // Y[r, 0] = 16 * 0.5 * (1 + 2 + 1 + 2) + 1 = 48 + 1 = 49. + // Y[r, 1] = 16 * 0.5 * 4 * (-2) + 2 = -64 + 2 = -62. + std::vector expected(m * n); + for (int64_t row = 0; row < m; ++row) { + expected[row * n + 0] = 49.0f; + expected[row * n + 1] = -62.0f; + } + + 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.AddOptionalInputEdge(); // input_scale + 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); +} + +// Scalar fallback: K % 8 == 0 but block_size is odd, so the two nibbles of a packed byte can land +// in different scale blocks. The vectorized kernel assumes they never do, so it must be skipped. +// Blocks are [0,5) [5,10) [10,15) [15,16) with scales 1.0, 2.0, 0.5, 1.0. +TEST(MatMulBlockQuantizedFp4WeightOpTest, PrefillDequantOddBlockSizeFp16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t m = 10; // > kGemvMaxM + constexpr int64_t n = 1; + constexpr int64_t k = 16; + constexpr int64_t block_size = 5; + constexpr int64_t k_blocks = 4; // ceil(16 / 5) + + std::vector b(n * (k / 2), 0x22); // all +1.0 + std::vector weight_scale = {0x38, 0x40, 0x30, 0x38}; + std::vector weight_scale_2 = {1.0f}; + + std::vector a(m * k, 1.0f); + // Y = 5 * 1.0 + 5 * 2.0 + 5 * 0.5 + 1 * 1.0 = 18.5. + std::vector expected(m * n, 18.5f); + + OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", block_size); + 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); +} + +// Scalar fallback: K % 8 != 0, so a thread cannot own an aligned 8-element chunk. +TEST(MatMulBlockQuantizedFp4WeightOpTest, PrefillDequantKNotMultipleOf8Bf16) { + if (!HasCudaEnvironment(800)) { + GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; + } + + constexpr int64_t m = 12; // > kGemvMaxM + constexpr int64_t n = 2; + constexpr int64_t k = 12; // 12 % 8 == 4 + constexpr int64_t block_size = 4; + constexpr int64_t k_blocks = k / block_size; + + // Row 0 = +1.0 (0x22), row 1 = +0.5 (0x11). + std::vector b(n * (k / 2)); + for (int64_t j = 0; j < k / 2; ++j) { + b[0 * (k / 2) + j] = 0x22; + b[1 * (k / 2) + j] = 0x11; + } + // Row 0 blocks {1.0, 2.0, 1.0}, row 1 all 1.0. + std::vector weight_scale = {0x38, 0x40, 0x38, 0x38, 0x38, 0x38}; + std::vector weight_scale_2 = {1.0f}; + + std::vector a(m * k, 1.0f); + // Y[r, 0] = 4 * (1*1 + 1*2 + 1*1) = 16. Y[r, 1] = 4 * 3 * 0.5 = 6. + std::vector expected(m * n); + for (int64_t row = 0; row < m; ++row) { + expected[row * n + 0] = 16.0f; + expected[row * n + 1] = 6.0f; + } + + OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); + test.AddAttribute("block_size", block_size); + 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 54fd1fc58f1d6e09dc1c4e9efa5cb28d745cb26a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 18 Aug 2026 00:00:58 +0000 Subject: [PATCH 2/2] fix(cuda): address fp4 dequant review feedback --- .../cuda/math/matmul_block_scaled_fp4.cu | 2 +- .../matmul_block_scaled_fp4_test.cc | 41 +++++++++++-------- 2 files changed, 25 insertions(+), 18 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 2143f7820b527..5b4cb416d8f6e 100644 --- a/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu +++ b/onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu @@ -276,7 +276,7 @@ __global__ void DequantizeNvFp4Vec8Kernel(T* __restrict__ out, const int left0 = pairs_per_block - (pair0 - blk0 * pairs_per_block); const float g = *weight_scale_2; - for (int row = blockIdx.y; row < n; row += gridDim.y) { + for (int64_t row = blockIdx.y; row < n; row += gridDim.y) { const uint32_t word = reinterpret_cast(b_packed + static_cast(row) * (k >> 1))[chunk]; const uint8_t* srow = weight_scale + static_cast(row) * k_blocks; 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 e366df6c34ade..0cdd7fd1d4fd1 100644 --- a/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc +++ b/onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc @@ -687,8 +687,8 @@ TEST(MatMulBlockQuantizedFp4WeightOpTest, GemvTensorCoreLaneOwnershipFp16) { // every one uses M > 8 so the decode GEMV is skipped and the dequant actually runs. // --------------------------------------------------------------------------- -// Vectorized path: K % 8 == 0, even block_size, several K blocks with different scales, and a -// negative weight row so the sign bit of the packed nibble is exercised. +// Vectorized path: K % 8 == 0, even block_size, scale boundaries inside every 8-element chunk, +// and a negative weight row so the sign bit of the packed nibble is exercised. TEST(MatMulBlockQuantizedFp4WeightOpTest, PrefillDequantVectorizedFp16) { if (!HasCudaEnvironment(800)) { GTEST_SKIP() << "CUDA device is required for MatMulBlockQuantizedFp4Weight."; @@ -696,8 +696,9 @@ TEST(MatMulBlockQuantizedFp4WeightOpTest, PrefillDequantVectorizedFp16) { constexpr int64_t m = 16; // > kGemvMaxM (8), so the dequant + cuBLAS path runs constexpr int64_t n = 3; - constexpr int64_t k = 32; - constexpr int64_t k_blocks = k / 16; + constexpr int64_t k = 24; + constexpr int64_t block_size = 6; + constexpr int64_t k_blocks = k / block_size; // Row 0 = +1.0 (0x22), row 1 = +2.0 (0x44), row 2 = -1.5 (sign 0x8 | 0x3 -> 0xBB). std::vector b(n * (k / 2)); @@ -706,30 +707,36 @@ TEST(MatMulBlockQuantizedFp4WeightOpTest, PrefillDequantVectorizedFp16) { b[1 * (k / 2) + j] = 0x44; b[2 * (k / 2) + j] = 0xBB; } - // Row 0 blocks {1.0, 2.0}, row 1 {1.0, 1.0}, row 2 {2.0, 0.5}. - std::vector weight_scale = {0x38, 0x40, 0x38, 0x38, 0x40, 0x30}; + // E4M3 scale bytes 0.5, 1.0, 2.0, and 4.0 arranged differently in each weight row. + std::vector weight_scale = { + 0x38, 0x40, 0x30, 0x48, + 0x40, 0x30, 0x48, 0x38, + 0x30, 0x48, 0x38, 0x40}; std::vector weight_scale_2 = {1.0f}; std::vector a(m * k); for (int64_t row = 0; row < m; ++row) { for (int64_t col = 0; col < k; ++col) { - a[row * k + col] = static_cast(row + 1); + a[row * k + col] = static_cast((row + 1) * (col + 1)); } } - // Each block contributes 16 * (row + 1) * value * block_scale: - // Y[r, 0] = 16 (r+1) (1*1 + 1*2) = 48 (r+1) - // Y[r, 1] = 16 (r+1) (2*1 + 2*1) = 64 (r+1) - // Y[r, 2] = 16 (r+1) (-1.5*2 + -1.5*0.5) = -60 (r+1) - std::vector expected(m * n); + constexpr float weight_values[] = {1.0f, 2.0f, -1.5f}; + constexpr float scale_values[][k_blocks] = { + {1.0f, 2.0f, 0.5f, 4.0f}, + {2.0f, 0.5f, 4.0f, 1.0f}, + {0.5f, 4.0f, 1.0f, 2.0f}}; + std::vector expected(m * n, 0.0f); for (int64_t row = 0; row < m; ++row) { - const float scale = static_cast(row + 1); - expected[row * n + 0] = 48.0f * scale; - expected[row * n + 1] = 64.0f * scale; - expected[row * n + 2] = -60.0f * scale; + for (int64_t col = 0; col < n; ++col) { + for (int64_t kk = 0; kk < k; ++kk) { + expected[row * n + col] += + a[row * k + kk] * weight_values[col] * scale_values[col][kk / block_size]; + } + } } OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain); - test.AddAttribute("block_size", 16); + test.AddAttribute("block_size", block_size); test.AddInput("A", {m, k}, FloatsToMLFloat16s(a)); test.AddInput("B", {n, k / 2}, b); test.AddInput("weight_scale", {n, k_blocks}, weight_scale);