onnxruntime

8b9258f1 - [CUDA] Bound the FP8 weight dequant scratch by tiling over N (#32129)

[CUDA] Bound the FP8 weight dequant scratch by tiling over N (#32129) ### Description The `MatMulBlockQuantizedFp8Weight` fallback path (`M > 8`, i.e. every prefill matmul) expands the packed weight into an `[N, K]` scratch buffer of the activation type and then calls cuBLAS. That scratch is **2.37 GiB** for a 248320 x 5120 LM head, and it stays live for the duration of the GEMM even though the GEMM reads it exactly once. This caps the scratch and runs the dequantize + GEMM pair over N tiles that fit inside the cap: ``` for n_offset in 0, tile_rows, 2 * tile_rows, ...: dequantize B[n_offset : n_offset + rows, :] into the scratch cublasGemmHelper(...) writing Y + n_offset ``` The row-major `[M, N]` output is column-major `[N, M]` to cuBLAS, so an N tile is a plain row offset into `Y` — there is no extra copy, no split-K reduction, and no change to the arithmetic performed per output element. `ORT_FP8_DEQUANT_SCRATCH_MIB` sets the cap in MiB (default 256). Shapes small enough to fit the cap take a single tile and are completely unaffected. ### Verification **Memory and speed** on Qwen3.8-27B at an 8K prompt (H200). The cap was chosen by sweeping it: | cap | peak memory | TTFT | |---|---|---| | untiled (before) | 32609 MiB | baseline | | 1 GiB | -3100 MiB | +5.2% | | **256 MiB (default)** | **28503 MiB (-4106)** | **+1.1%** | | 128 MiB | -4106 MiB | +13.9% | **Numerics.** Per-element arithmetic is unchanged, but tiling changes the N extent handed to cuBLAS, so the library is free to select a different kernel for a tile than for the whole matrix. Measured tiled vs untiled, FP16, `K = 5120`, `block_size = 128`: | M | N | tiles | max abs diff | differing elements | |---:|---:|---:|---:|---:| | 16 | 6144 | 1 | 0 | 0 | | 1024 | 6144 | 1 | 0 | 0 | | 16 | 32768 | 2 | 0.125 | 28065 | | 1024 | 32768 | 2 | 0.125 | 2396209 | | 64 | 248320 | 10 | 0 | 0 | `0.125` is exactly one FP16 ULP at the magnitude of these outputs (values around 200, where FP16 spacing is `2^-3`). So where the split changes cuBLAS's kernel choice the result moves by at most one last-bit rounding step — the same order of change as cuBLAS picking a different tactic for any other reason. Single-tile shapes are bitwise identical. **Tests.** `WeightDequantScratchTilingFp16` sets `ORT_FP8_DEQUANT_SCRATCH_MIB=1` so a small shape (`N = 769`, `K = 4096`) still splits into six tiles, with a distinct scale on every N row so that a tile reading the wrong weight/scale offset or writing the wrong output column changes `Y`. Confirmed under nsys that the test issues 3 dequant launches while a single-tile control issues 1. All 11 `MatMulBlockQuantizedFp8WeightOpTest` cases pass. ### Motivation and Context On Qwen3.8-27B the untiled scratch was the single largest transient allocation in the model and pushed peak memory ~4.1 GB above what the weights and KV cache actually need. Recovering it makes room for longer contexts or a larger batch at the same memory budget, and costs ~1% of TTFT. This is independent of #32128 (which speeds up the NVFP4 dequantization kernel); the two touch different operators and can land in either order.