onnxruntime

2d211f8c - [CUDA] Vectorize the NVFP4 weight dequantization for prefill (#32128)

[CUDA] Vectorize the NVFP4 weight dequantization for prefill (#32128) ### Description `MatMulBlockQuantizedFp4Weight` falls back to "expand the packed weight into an `[N, K]` scratch buffer, then cuBLAS" whenever the decode GEMV and the native SM120 path do not apply — i.e. for every prefill matmul on Hopper. That expansion was the single most expensive kernel in NVFP4 prefill. The old `DequantizeNvFp4Kernel` gave each thread one packed byte (2 FP4 codes): | defect | cost | |---|---| | 1-byte load + two separate 2-byte stores | no vectorization | | `idx / half_k` and `k0 / block_size` | two integer divisions per thread | | `*weight_scale_2` read per thread | a global load per thread | | `__nv_cvt_fp4x2_to_halfraw2()` | software-emulated on pre-Blackwell: the SASS has branches **and** a subnormal normalization loop | `DequantizeNvFp4Vec8Kernel` replaces it when `K % 8 == 0` and `block_size` is even. 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. The row index comes from `blockIdx.y`, `weight_scale_2` is hoisted into a register, and the scale index advances incrementally instead of by per-element division. Codes are decoded with the branch-free `Fp4Cvt` `prmt` lookup already in this file (added for the decode GEMV in #31155). The scalar kernel is kept for odd `block_size` or `K % 8 != 0`. One design point worth recording: **8 elements per thread, not more.** Widening the per-thread chunk to 32 elements (four back-to-back `uint4` stores) makes every store instruction stride across lanes. Against a copy-only kernel with identical index math, that shape ceilings at 1.9 TB/s while one `uint4` store per thread reaches 3.9 TB/s on H200 — a 2x difference that no amount of tile tuning recovers. ### Verification **Bitwise identical.** Dequantization is elementwise and `Fp4Cvt` reproduces the intrinsic's bit pattern exactly, so no output should change. Checked three ways: - `Fp4Cvt` vs `__nv_cvt_fp4x2_to_halfraw2()` brute-forced over **all 256 packed byte values** — 0 mismatches (including `-0.0` for code `0x8`). - Op-level SHA-256 over the full output for 8 shapes (both dtypes, `block_size` 16/32, `K % 32 == 0`, `K % 32 == 16`, `K % 8 == 4`, `N`/`K` from 128 to 5120) — identical before and after. - Qwen3.8-27B NVFP4, 8K prompt / 128 generated / MTP `N=3` on H200: the generated-token SHA-256 is unchanged (`095222fc5fdb…`) across 3 runs per arm. **Kernel time** (H200, `M = 1024`, BF16, `block_size = 16`, median over 50 iterations): | 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 | **Tests.** Four new cases in `matmul_block_scaled_fp4_test.cc`, all with `M > 8` so the decode GEMV is skipped and the dequant actually runs. Existing FP4 tests all used `M <= 8`, so the prefill path had no coverage at prefill shapes. Each case was confirmed under nsys to reach the intended kernel: | test | kernel reached | |---|---| | `PrefillDequantVectorizedFp16` | `DequantizeNvFp4Vec8Kernel<__half>` | | `PrefillDequantVectorizedBiasBf16` | `DequantizeNvFp4Vec8Kernel<__nv_bfloat16>` | | `PrefillDequantOddBlockSizeFp16` | `DequantizeNvFp4Kernel<__half>` | | `PrefillDequantKNotMultipleOf8Bf16` | `DequantizeNvFp4Kernel<__nv_bfloat16>` | `PrefillDequantOddBlockSizeFp16` is the interesting one: with an odd `block_size` the two nibbles of a packed byte can land in different scale blocks, which is exactly the assumption the vectorized kernel makes and therefore the reason it must be skipped. All 17 `MatMulBlockQuantizedFp4WeightOpTest` cases pass. ### Motivation and Context Measured on Qwen3.8-27B NVFP4 (168 `MatMulBlockQuantizedFp4Weight` nodes holding 14.97 G weights), 8K prompt, H200: - `DequantizeNvFp4Kernel` was **46.9% of all prefill GPU time** (1750 ms of 3732 ms) — more than every cuBLAS GEMM in the model combined. - After this change it is **17.1% (409 ms)**, and total prefill GPU time drops 3732 -> 2394 ms. - End-to-end TTFT for 8K/128/spec-3: **3876 -> 2618 ms (-32.5%)**, averaged over 3 interleaved runs per arm. Decode throughput and MTP acceptance are unchanged, as expected — decode uses the GEMV path and never reaches this kernel. The kernel is now at the bandwidth bound for the work it does: one full dequantization pass over these weights moves 34.86 GiB, which is 12.5 ms at ~3 TB/s, and the measured cost is 12.4 ms per pass.