[CUDA] Vectorize the NVFP4 weight dequantization for prefill

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

onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp4.cu

@@ -241,6 +241,72 @@ struct Fp4Cvt<nv_bfloat16> {
   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 <typename T>
+__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<T>;
+  using T2 = typename Cvt::T2;
+
+  const int chunk = static_cast<int>(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 (int64_t row = blockIdx.y; row < n; row += gridDim.y) {
+    const uint32_t word = reinterpret_cast<const uint32_t*>(b_packed + static_cast<size_t>(row) * (k >> 1))[chunk];
+    const uint8_t* srow = weight_scale + static_cast<size_t>(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<T>(v.x * s);
+      outv[2 * j + 1] = from_float<T>(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<uint4*>(out + static_cast<size_t>(row) * k)[chunk] = *reinterpret_cast<const uint4*>(outv);
+  }
+}
+
 template <typename T, int RowsPerBlock>
 __global__ __launch_bounds__(32 * kGemvWarpsPerBlock,
                              GemvMinBlocksPerSm<RowsPerBlock>::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<int>((total + kThreads - 1) / kThreads);
   const uint8_t* bp = reinterpret_cast<const uint8_t*>(b_packed);
   const uint8_t* ws = reinterpret_cast<const uint8_t*>(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<unsigned int>((chunks_per_row + threads - 1) / threads);
+    const unsigned int grid_y = static_cast<unsigned int>(n < 65535 ? n : 65535);
+    const dim3 blocks{grid_x, grid_y};
+    if (is_bf16) {
+      DequantizeNvFp4Vec8Kernel<nv_bfloat16><<<blocks, threads, 0, stream>>>(
+          reinterpret_cast<nv_bfloat16*>(b_dequant), bp, ws, weight_scale_2, n, k, k_blocks, block_size,
+          chunks_per_row);
+    } else {
+      DequantizeNvFp4Vec8Kernel<half><<<blocks, threads, 0, stream>>>(
+          reinterpret_cast<half*>(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<int>((total + kThreads - 1) / kThreads);
+
   if (is_bf16) {
     DequantizeNvFp4Kernel<nv_bfloat16><<<blocks, kThreads, 0, stream>>>(
         reinterpret_cast<nv_bfloat16*>(b_dequant), bp, ws, weight_scale_2, n, k, k_blocks, block_size);

onnxruntime/test/contrib_ops/matmul_block_scaled_fp4_test.cc

@@ -679,6 +679,204 @@ 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, 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.";
+  }
+
+  constexpr int64_t m = 16;  // > kGemvMaxM (8), so the dequant + cuBLAS path runs
+  constexpr int64_t n = 3;
+  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<uint8_t> 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;
+  }
+  // E4M3 scale bytes 0.5, 1.0, 2.0, and 4.0 arranged differently in each weight row.
+  std::vector<uint8_t> weight_scale = {
+      0x38, 0x40, 0x30, 0x48,
+      0x40, 0x30, 0x48, 0x38,
+      0x30, 0x48, 0x38, 0x40};
+  std::vector<float> weight_scale_2 = {1.0f};
+
+  std::vector<float> 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<float>((row + 1) * (col + 1));
+    }
+  }
+  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<float> expected(m * n, 0.0f);
+  for (int64_t row = 0; row < m; ++row) {
+    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<int64_t>("block_size", block_size);
+  test.AddInput<MLFloat16>("A", {m, k}, FloatsToMLFloat16s(a));
+  test.AddInput<uint8_t>("B", {n, k / 2}, b);
+  test.AddInput<uint8_t>("weight_scale", {n, k_blocks}, weight_scale);
+  test.AddInput<float>("weight_scale_2", {1}, weight_scale_2);
+  test.AddOutput<MLFloat16>("Y", {m, n}, FloatsToMLFloat16s(expected));
+  test.SetOutputTolerance(0.5f);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<uint8_t> 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<uint8_t> weight_scale = {0x38, 0x40, 0x38, 0x40, 0x38, 0x38, 0x38, 0x38};
+  std::vector<float> weight_scale_2 = {0.5f};
+  std::vector<float> bias = {1.0f, 2.0f};
+
+  std::vector<float> 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<float> 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<int64_t>("block_size", 16);
+  test.AddInput<BFloat16>("A", {m, k}, FloatsToBFloat16s(a));
+  test.AddInput<uint8_t>("B", {n, k / 2}, b);
+  test.AddInput<uint8_t>("weight_scale", {n, k_blocks}, weight_scale);
+  test.AddInput<float>("weight_scale_2", {1}, weight_scale_2);
+  test.AddOptionalInputEdge<float>();  // input_scale
+  test.AddInput<BFloat16>("bias", {n}, FloatsToBFloat16s(bias));
+  test.AddOutput<BFloat16>("Y", {m, n}, FloatsToBFloat16s(expected));
+  test.SetOutputTolerance(0.5f);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<uint8_t> b(n * (k / 2), 0x22);  // all +1.0
+  std::vector<uint8_t> weight_scale = {0x38, 0x40, 0x30, 0x38};
+  std::vector<float> weight_scale_2 = {1.0f};
+
+  std::vector<float> a(m * k, 1.0f);
+  // Y = 5 * 1.0 + 5 * 2.0 + 5 * 0.5 + 1 * 1.0 = 18.5.
+  std::vector<float> expected(m * n, 18.5f);
+
+  OpTester test("MatMulBlockQuantizedFp4Weight", 1, onnxruntime::kMSDomain);
+  test.AddAttribute<int64_t>("block_size", block_size);
+  test.AddInput<MLFloat16>("A", {m, k}, FloatsToMLFloat16s(a));
+  test.AddInput<uint8_t>("B", {n, k / 2}, b);
+  test.AddInput<uint8_t>("weight_scale", {n, k_blocks}, weight_scale);
+  test.AddInput<float>("weight_scale_2", {1}, weight_scale_2);
+  test.AddOutput<MLFloat16>("Y", {m, n}, FloatsToMLFloat16s(expected));
+  test.SetOutputTolerance(0.5f);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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<uint8_t> 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<uint8_t> weight_scale = {0x38, 0x40, 0x38, 0x38, 0x38, 0x38};
+  std::vector<float> weight_scale_2 = {1.0f};
+
+  std::vector<float> 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<float> 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<int64_t>("block_size", block_size);
+  test.AddInput<BFloat16>("A", {m, k}, FloatsToBFloat16s(a));
+  test.AddInput<uint8_t>("B", {n, k / 2}, b);
+  test.AddInput<uint8_t>("weight_scale", {n, k_blocks}, weight_scale);
+  test.AddInput<float>("weight_scale_2", {1}, weight_scale_2);
+  test.AddOutput<BFloat16>("Y", {m, n}, FloatsToBFloat16s(expected));
+  test.SetOutputTolerance(0.5f);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> 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