onnxruntime
e8ae6ce9 - Add float zero point support for 2-bit LUT GEMM in MatMulNBits (#28354)

Commit
113 days ago
Add float zero point support for 2-bit LUT GEMM in MatMulNBits (#28354) ## Description Adds support for float/float16 zero points in the 2-bit MatMulNBits LUT GEMM path, enabling AMD QAD/Quark 2-bit quantization which requires a fractional zero point of 1.5. Addresses #28162 ### Problem QAD 2-bit quantization uses non-uniform levels `[-1, -1/3, 1/3, 1]`, expressed via `dequant = (q - 1.5) * scale`. The zero point 1.5 cannot be represented as a packed uint8 value. The existing LUT GEMM packing API only accepted `uint8_t*` zero points, and the fallback dequant path crashed with `ORT_ENFORCE(nbits_ == 4)` when encountering 2-bit + float ZP. ### Changes **MLAS layer** — Widened `MlasLutGemmPack()` to accept `const void* QuantBZeroPoint` + `bool IsFloatZeroPoint`, following the existing `MlasQNBitGemmPackQuantBData` convention. The AVX2 packer reads float ZP values directly per quantization group when `IsFloatZeroPoint` is set, computing the same `(zp - midpoint) * scale` correction stored in the packed buffer. The compute kernel (`TMACComputeGemm_avx2`) is unchanged — it already consumes ZP as a float correction during accumulation. **MatMulNBits CPU kernel** — Relaxed the PrePack early-exit guard to allow float ZP into the LUT GEMM path (not non-LUT paths). Added fp16→fp32 conversion for ZP tensors, matching how scales are already handled. Fixed the Compute() path to null out prepacked zero_points to avoid a null dereference in CheckInputs. Fixed the 2-bit fallback dequant path: relaxed the `nbits_==4` enforce, added inline 2-bit scalar dequant for float and MLFloat16 ZP with correct packed-B indexing for padded K shapes. **Tests** — Added MLAS-level float ZP tests across block lengths 32/64/128 with ZP values {0, 1.5, 2, 3}. Added provider-level directed QAD tests (`zp=1.5`) verifying end-to-end correctness through the LUT GEMM path. ### Testing - 72 MLAS LUT GEMM tests pass (including 36 new float ZP tests) - 13 provider-level 2-bit tests pass (including new QAD float ZP tests) - No regressions in existing uint8 ZP tests - lintrunner clean ### Files changed | File | Change | |------|--------| | `core/mlas/inc/mlas_qnbit.h` | API: `void*` ZP + `IsFloatZeroPoint` flag | | `core/mlas/lib/qlutgemm.h` | Dispatch typedef update | | `core/mlas/lib/qlutgemm.cpp` | Pass-through plumbing | | `core/mlas/lib/sqnbitgemm_lut_kernel_avx2.cpp` | Float ZP packing branch | | `contrib_ops/cpu/quantization/matmul_nbits.cc` | PrePack guard, fallback fix, ZP validation | | `test/mlas/unittest/test_sqlutgemm.cpp` | Float ZP MLAS tests | | `test/mlas/bench/bench_lutgemm.cpp` | Updated call signature | | `test/contrib_ops/matmul_2bits_test.cc` | Float ZP provider tests | --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading