onnxruntime

0a855f8e - [MLAS] Run fp16 MatMulNBits through the int8 path without fp32 temporaries (#29791)

[MLAS] Run fp16 MatMulNBits through the int8 path without fp32 temporaries (#29791) ### Description For a fp16 model at accuracy_level 4, MatMulNBits runs the int8 (CompInt8) path. Since #27820 that path is reached by converting the whole activation matrix A from fp16 to fp32, running the int8 GEMM in fp32, and converting the whole result back to fp16. That is two full M x K and M x N fp32 temporaries allocated and streamed on every Compute call, on top of the GEMM itself. This removes both. On x64 the int8 path now: - quantizes A straight from fp16 to int8 in the workspace step (the fp16 elements are widened to float with F16C inside the existing quantize pass), so there is no fp32 copy of A, and - writes the output as fp16 as the kernel produces it. Each worker runs its GEMM strip into a per-thread fp32 scratch and converts it to fp16 in place, so the full fp32 result is never materialized and there is no separate conversion pass. The scratch is sized to the strip the worker computes: at most 128 x 128 floats on the threaded path, M x 128 on the single-threaded path, where the whole row range goes through one kernel call to keep the result bit-identical to the fp32 path. The scale and bias are still pre-converted to fp32 once at prepack, as before. Nothing about the GEMM math changes, so the output is bit-identical to the current path: the quantized A is the same bytes the fp32 conversion would produce, and the fp16 output is the same bytes you get by converting the fp32 result with MlasConvertFloatToHalfBuffer. The change is guarded so the fp32 (SQNBIT_CompFp32) path and the fp32-input int8 path are untouched; it only triggers for fp16 input on the int8 path where the platform provides the fp16 quantizer, which today is the x64 AVX2 and AVX2-VNNI dispatches. Machines that select the AVX-512 dispatch, and other targets, keep the existing behavior. Numbers from a Core Ultra 5 225 (Arrow Lake, AVX2 + AVX-VNNI), 4-bit weights, accuracy_level 4, default thread count, min of best-of-150 over interleaved rounds. The fp32 column is the same int8 GEMM with fp32 input and is unchanged by this PR; it is the floor the fp16 path is trying to reach. Prefill, fp16 latency in ms (min of best-of-150 across 8 interleaved rounds): | M x K x N | fp16 before | fp16 after | fp32 | fp16 before/fp32 | fp16 after/fp32 | |---|---|---|---|---|---| | 128 x 4096 x 4096 | 4.36 | 3.36 | 3.30 | 1.32 | 1.02 | | 512 x 4096 x 4096 | 20.46 | 16.39 | 16.33 | 1.25 | 1.00 | | 512 x 11008 x 4096 | 53.29 | 44.98 | 44.33 | 1.20 | 1.01 | | 1024 x 4096 x 4096 | 42.69 | 33.99 | 33.41 | 1.28 | 1.02 | fp16 prefill was 20 to 32% slower than the same GEMM with fp32 input; with the two temporaries gone it lands within about 2% of fp32. That is roughly a 20% cut in fp16 prefill latency. Decode (M=1) is unchanged either way, as expected, since its activation and output are a single row. The single-thread numbers move the same direction (fp16 comes down to fp32) at a smaller before-gap, since the conversions are not parallel there to begin with. End to end on a real model, a Phi-3-mini int4 model run in fp16 does a 512-token prefill about 18% faster with this change (min of best over 8 interleaved rounds on one machine). The fp16 model lands within about 4% of the same model run in fp32, where before it was around 23% behind. That is a whole-model prefill number so it sits a little under the per-op matmul figure above, which is expected since the rest of the model is unchanged. Testing: - new test_sqnbitgemm_fp16_quant_a.cpp checks the fp16 paths are exact: the GEMM output from fp16 A is bit-identical to the output from the same A converted to fp32 first, for 2-, 4- and 8-bit weights, and the 4-bit direct fp16 output equals the fp32 result converted to fp16, over block sizes 16-256, symmetric and asymmetric, with and without bias, single and multi threaded. - the existing MatMulNBits fp16 operator tests pass unchanged. - full MLAS unit test suite green. ### Motivation and Context #27251 asked for a faster fp16 MatMulNBits on CPU. #27820 got most of the way by routing fp16 through the fp32 int8 path, and noted the conversion overhead it introduced (the 8-bit accuracy_level 4 regression in that PR is part of it). This takes the rest: with the temporaries gone, the fp16 int8 path does essentially the same work as the fp32 int8 path. The win is largest at prefill, where the temporaries are large; decode is unaffected since its activation and output are tiny. On x64 the A quantizer runs in the workspace step shared by every CompInt8 bit width, so it covers 2-, 4- and 8-bit weights; the fp16 output epilogue is in the 4-bit CompInt8 kernel. The 8-bit output epilogue and an AVX-512 A quantizer are natural follow-ups. Happy to split this into the A side and the C side as two PRs if that is easier to review.