Add float zero point support for 2-bit LUT GEMM in MatMulNBits

[vraspar]

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

[Copilot]

Pull request overview

Adds float/float16 zero-point handling for the CPU 2-bit MatMulNBits path, primarily by widening the MLAS LUT pack API and teaching the CPU kernel to route/convert float zero points for LUT prepack and unpacked fallback dequantization.

Changes:

• Extends MLAS LUT GEMM packing APIs and AVX2 packing logic to accept float zero-point data.
• Updates the CPU MatMulNBits kernel to allow LUT prepack with unquantized zero points and adds 2-bit float/FP16 fallback dequantization.
• Adds MLAS/unit, provider, and benchmark call-site updates for the new LUT pack signature and float-ZP scenarios.

Findings

• onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc: relaxing the early exit for unquantized zero points now lets the LUT prepack path run even when zero_points is dynamic. Because LUT packing only uses TryGetConstantInput and has no input_idx == zero_points pack step, dynamic float/float16 zero points will be ignored in the packed path.
• onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc: both new zero-point conversion blocks assume every non-float zero-point tensor is MLFloat16. The schema also allows bfloat16, so constant BF16 zero points would be reinterpreted as FP16 during prepack.
• onnxruntime/test/contrib_ops/matmul_2bits_test.cc: the new provider test never checks MlasIsLutGemmAvailable(), so on platforms without LUT GEMM it can pass via the unpacked fallback and fail to validate the new LUT float-ZP path.
• onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc: the new float zero-point 2-bit fallback dequant branch is not covered by the added tests, because the new tests force LUT GEMM on LUT-compatible shapes.
• onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc: the new MLFloat16-specific 2-bit float zero-point fallback branch also lacks coverage; the added tests only exercise float inputs/zero points.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_sqlutgemm.cpp	Adds MLAS float-zero-point LUT GEMM unit coverage and new pack call signature usage.
onnxruntime/test/mlas/bench/bench_lutgemm.cpp	Updates benchmark call sites for the widened LUT pack API.
onnxruntime/test/contrib_ops/matmul_2bits_test.cc	Adds provider-level float zero-point 2-bit tests.
onnxruntime/core/mlas/lib/sqnbitgemm_lut_kernel_avx2.cpp	Implements AVX2 packing support for float zero points.
onnxruntime/core/mlas/lib/qlutgemm.h	Updates LUT dispatch typedef to carry generic ZP pointer + type flag.
onnxruntime/core/mlas/lib/qlutgemm.cpp	Threads the new zero-point arguments through MLAS LUT pack plumbing.
onnxruntime/core/mlas/inc/mlas_qnbit.h	Updates public MLAS LUT pack declaration/docs for float zero points.
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc	Adjusts CPU prepack/compute logic for float zero points in LUT and fallback paths.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

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

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[vraspar]

Review: Add float zero point support for 2-bit LUT GEMM in MatMulNBits

Overall: Clean, well-scoped PR. The math is correct, the API extension is pragmatic, and test coverage is solid. Two items worth addressing:

---

1. Test reference uses wrong packed stride (both test files)

matmul_2bits_test.cc -- RunTest2BitsFloatZP:
cpp int64_t packed_idx = n * (K / 4) + k / 4; // assumes K == packed_k

test_sqlutgemm.cpp -- TestFloatZeroPoint:
cpp size_t packed_idx = n * (K / (8 / BlkBitWidth)) + k / (8 / BlkBitWidth); // same issue

MlasQuantizeBlockwise pads to packed_k = k_blocks * block_size. The kernel correctly uses the padded stride (packed_k / 4), and the fallback test (Float32_2b_FloatZP_Fallback) gets it right too. These two references diverge from the kernel when K is not a multiple of block_size. All current test cases happen to use aligned K, so it is not a live failure -- but it would produce wrong reference values if someone adds RunTest2BitsFloatZP(1, 128, 100, 32, 1.5f). Suggest using the same packed_k-based stride and adding one non-aligned-K test case to prevent regression.

2. Duplicated ZP conversion in PrePack -- extract helper

The ~25-line fp16-to-fp32 ZP conversion + validation block is copy-pasted in two PrePack code paths (input_idx == B and input_idx == scales under prefer_lut_gemm_). If bfloat16 ZP or other changes come later, both copies need updating. Consider extracting into a small helper method or local lambda.

[tianleiwu]

Review Summary

The approach is sound: widening the MLAS API to accept void* ZP with a type flag, scalar fallback for 2-bit float ZP, and properly nulling zero_points in Compute() when LUT GEMM has consumed them. The ConvertFloatZeroPointsForLutGemm helper and the compute_zp_correction lambda in the AVX2 kernel are clean abstractions.

However, the cross-type ZP issue from the previous review round remains unaddressed: in ComputeBUnpacked<float>, an MLFloat16 zero-point tensor enters the uint8 cast path, producing garbage. Additionally, there's still no test coverage for the MLFloat16 ZP paths that this PR adds code for.

Overall quality is good — well-structured, good test coverage of the primary path, and clear commit messages explaining the progression of fixes.

[tianleiwu]

More inline comments...

[vraspar]

Addressing review feedback (commit 004beb1)

Pushed fixes for all outstanding review items. Summary:

Fixes applied

Change	Addresses
K % BlkLen == 0 guard in MlasIsLutGemmAvailable	Thread by @tianleiwu re stride consistency — the AVX2 packer uses floor division (K / BlkLen) which diverges from caller's ceiling division when K % BlkLen != 0. This is pre-existing (affects uint8 ZP too) but newly exposed. Guard is the safe minimal fix; full arbitrary-K support is a separate MLAS effort.
Dynamic ZP guard in PrePack	Own finding — if ZP is a non-constant input, has_zp_input_ is false, LUT packs without ZP, and Compute() nulls ZP → silent wrong results. Added has_zp_arg_ member + early return, similar to KleidiAI's dynamic ZP fallback (lines 377-386).
int32_t → size_t in fallback dequant loops	Thread 18 by @tianleiwu — consistent with MLAS conventions, prevents overflow for very large models.
T3 constraint comments at ZP dispatch points	Thread 17 by @tianleiwu — the cross-type ZP scenario (MLFloat16 ZP with float activations) cannot happen because the kernel registration constrains T3 to {uint8_t, T1}. Added comments at the dispatch points explaining this so future readers don't need to trace through kernel binding.
MLFloat16 ZP fallback test	Thread 20 by @tianleiwu — new test with MLFloat16 activations + MLFloat16 zero points exercising ComputeBUnpacked<MLFloat16> 2-bit dequant path.

Items deferred (with rationale)

• Full arbitrary-K LUT GEMM support: Requires fixing packer, sizing, and compute kernel stride math across MLAS. The K % BlkLen guard prevents the bug; no known model needs non-aligned K with LUT GEMM.
• MLFloat16 LUT GEMM support: MatMulNBits<MLFloat16>::PrePack exits early for unquantized ZP. Adding LUT support for fp16 activations is a feature request with no current user demand; the fallback dequant path handles it correctly.

[tianleiwu]

Thanks for the follow-up updates. I rechecked the current head and did not find a new blocking correctness issue. I left two suggestion-level test coverage notes that would make the float zero-point changes harder to regress.

[tianleiwu]

Review Summary

The implementation is correct and well-structured. No correctness, security, or performance issues found.

Highlights:

• MLAS API widening — const void* QuantBZeroPoint + bool IsFloatZeroPoint is clean type erasure. All callers updated consistently.
• K % BlkLen guard — The new MlasIsLutGemmAvailable() check that rejects K % BlkLen != 0 closes a subtle floor/ceiling division mismatch between the caller and the AVX2 packer kernel.
• PrePack guard logic — The split into three distinct checks (has_g_idx_, has_unquantized_zero_point_ && !prefer_lut_gemm_, prefer_lut_gemm_ && has_zp_arg_ && !has_zp_input_) is correct. The dynamic ZP guard prevents silent zero-point ignorance when ZP is not a constant initializer.
• AVX2 compile-time dispatch — <bool HasZeroPoint, bool IsFloatZeroPoint> template with if constexpr ensures zero runtime overhead in the tight packing loop.
• Compute path safety — Nulling zero_points when prefer_lut_gemm_ && packed_b_ prevents null dereference in CheckInputs. When the dynamic ZP guard fires (no prepack), packed_b_ stays null and ZP is correctly fetched from context for the scalar fallback.
• Test coverage — Comprehensive: uniform ZP (QAD zp=1.5), varying per-block ZP (detects stride/transpose bugs), dynamic ZP fallback, non-aligned K, MLFloat16 ZP, MLAS-level tests across block lengths 32/64/128.

All four previously unresolved review threads have been addressed by the current head and resolved:

• Cross-typed ZP concern → correct per kernel registration constraints (T3 = {uint8_t, T1})
• Uniform ZP test concern → addressed by VaryingPerBlock test
• Dynamic ZP test concern → addressed by DynamicZP test with is_initializer=false
• MLFloat16 test concern → addressed by MLFloat16_2b_MLFloat16ZP_Fallback test