[Tianlei Wu (tianleiwu)]

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.

[Copilot]

Pull request overview

Caps FP8 weight-dequantization scratch memory by tiling GEMM work over the N dimension.

Changes:

• Adds configurable N-dimension tiling with a 256 MiB default cap.
• Adds FP16 tiling coverage.
• Documents behavior and performance results.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 3 comments.

File	Description
onnxruntime/contrib_ops/cuda/math/matmul_block_scaled_fp8.cc	Implements tiled dequantization and GEMM.
onnxruntime/test/contrib_ops/matmul_block_scaled_fp8_test.cc	Adds tiling correctness coverage.
docs/contrib_ops/cuda/matmul_block_scaled_fp8.md	Documents configuration and benchmarks.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

/azp run Linux Android Emulator QNN CI Pipeline

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Baiju Meswani (baijumeswani)]

PR description still says N=384, three tiles, and 10 tests. The current test uses N=769, seven tiles, and brings the suite to 11 tests.