[MLAS] Integrate Q4 asymmetric KleidiAI kernels into MatMulNBits

[martin-klacer-arm]

Description

This PR integrates KleidiAI asymmetric Q4 kernels into the MatMulNBits path for SQNBIT_CompInt8 when zero points are available during prepack.

The operator-facing changes are limited to the KleidiAI-eligible packed path: MatMulNBits now passes constant zero points into MLAS prepacking, marks shared prepacked buffers as finalized only when scales/zero points are actually folded into the packed B buffer, and declines prepacking for runtime zero-point inputs for the KleidiAI path so those cases keep using the existing fallback behavior.

Within MLAS, this adds asymmetric Q4 backend selection, RHS/LHS packing for the KleidiAI asymmetric packed layout, and dispatch for the corresponding GEMV/GEMM kernels. Cases that can't use the native asymmetric path continue to use the existing symmetric packed path and zero-point correction behavior.

Motivation and Context

Asymmetric block-quantized Q4 previously didn't have a dedicated execution path using KleidiAI - it was routed through the shared symmetric Q4 path with an additional zero-point correction step.

This PR adds a path that processes asymmetric zero points directly, allowing supported prepacked asymmetric weights to use KleidiAI kernels that handle zero points natively.

The performance was tested on a Mac Mini with M4 Pro using clean Release builds and onnxruntime_perf_test -m times -x 1 -s -S 1. The tables report the mean across 5 outer benchmark command invocations, each running 20 timed inferences for the E2E model (100 total) and 1000 timed inferences for synthetic nodes (5000 total).

Most of the uplift provided is thanks to the newly integrated asymmetric kernels using SME2 if available, and the asymmetric/symmetric split leaves room for future integration into either path.

The primary performance uplift was observed on an end-to-end test with the Phi 3.5 Mini Instruct ONNX model, using the CPU Int4 AWQ Block-128 variant.

Model	main mean ms	branch mean ms	main/branch uplift	timed inference repeats
Phi 3.5 AWQ	45.774	24.901	1.84	100

Additional uplift was observed on synthetic MatMulNBits cases using various shapes derived from Qwen 2.5 1.5B Instruct AWQ Block-128 Int4 asymmetric.

Shape	main mean ms	branch mean ms	main/branch uplift	timed inference repeats
M=1, N=1536, K=1536	0.028	0.015	1.91	5000
M=16, N=1536, K=1536	0.208	0.084	2.47	5000
M=128, N=1536, K=1536	1.665	0.649	2.57	5000
M=1, N=2048, K=1536	0.037	0.019	2.01	5000
M=2, N=2048, K=1536	0.089	0.033	2.72	5000
M=16, N=2048, K=1536	0.276	0.107	2.57	5000
M=128, N=2048, K=1536	2.458	0.831	2.96	5000
M=1, N=2047, K=1536	0.038	0.019	2.01	5000
M=16, N=2047, K=1536	0.275	0.107	2.56	5000

In addition, the measured symmetric control model exercising the existing symmetric path was unchanged as expected.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

This PR extends the MLAS/KleidiAI-backed MatMulNBits CPU path to support asymmetric Q4 (Int4) weights for SQNBIT_CompInt8 by selecting new KleidiAI asymmetric kernels and ensuring zero-points are handled correctly during prepack and shared-prepack reuse.

Changes:

• Adds KleidiAI asymmetric Q4 backend selection/packing in MLAS, plus dispatch support (alignment + “needs ZP correction” capability queries).
• Updates MatMulNBits prepack/sharing behavior to (a) pass constant zero-points into packing when available, and (b) decline prepack for runtime zero-points on the KleidiAI packed path to avoid incorrect sharing.
• Adds targeted unit tests for KleidiAI asymmetric Q4 and for dynamic zero-point + shared-prepack negative cases.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_sqnbitgemm.cpp	Plumbs backend selector config through test harness; adds KleidiAI asymmetric Q4 CompInt8 coverage.
onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.h	Adds a sharing-mode that expects prepack to be declined.
onnxruntime/test/contrib_ops/matmul_nbits_prepack_sharing_test_util.cc	Implements “expect no prepack” assertions for the new sharing-mode.
onnxruntime/test/contrib_ops/matmul_4bits_test.cc	Adds dynamic zero-point tests and shared-prepack negative-control coverage for KleidiAI asymmetric Q4 CompInt8.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8.cpp	Adds asymmetric backend selection and packed-kernel dispatch plumbing for ARM64/KleidiAI Q4.
onnxruntime/core/mlas/lib/qnbitgemm.h	Extends MLAS QNBit dispatch ABI for packed Q4 CompInt8 (alignment + ZP-correction capability + kernel config).
onnxruntime/core/mlas/lib/qnbitgemm.cpp	Uses new dispatch hooks to select alignment/correction behavior and pass backend config through packed paths.
onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h	Declares new ARM64/KleidiAI Q4 backend selection helpers and packed-path hooks.
onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp	Implements asymmetric Q4 RHS packing and backend selection for KleidiAI; wires new dispatch pointers.
onnxruntime/core/mlas/lib/kai_ukernel_interface.h	Adds asymmetric Q4 ukernel wrapper types/layout metadata and accessors.
onnxruntime/core/mlas/lib/kai_ukernel_interface.cpp	Registers/selects KleidiAI asymmetric Q4 ukernels (dot/i8mm/SME2 variants) and exposes accessors.
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc	Updates prepack/sharing logic for constant-vs-runtime zero-points and for “finalized packed B” semantics on the KleidiAI packed path.

[martin-klacer-arm]

Thank you for the review comments. They seem sensible, I'll upload a patch with fixes as soon as it's ready.

[martin-klacer-arm]

Thank you for the comments. I addressed the 'params' concerns as suggested by zero-initializing the variables. To address the memory concerns, I replaced the full-N temporary buffers with a reused nr-wide panel scratch buffer. Measuring with malloc_history -allEvents shows an approx. 98% reduction in cumulative scratch allocation for this specific part of the session creation and prepack.

[hariharans29]

Review: PR #29654 — [MLAS] Integrate Q4 asymmetric KleidiAI kernels into MatMulNBits

Approve. Introduces a dedicated KleidiAI backend (Qsi8d32pQai4c32p) for asymmetric block-quantized Q4 that consumes RHS zero points natively, instead of routing them through the symmetric packed path with a post-hoc C += AFloatBlkSum * BZpCorr^T correction. The main design win is running the SME2/i8mm/dotprod asymmetric kernels directly (Phi-3.5 Mini AWQ Block-128 Int4: 45.77 → 24.90 ms end-to-end on M4 Pro, 1.84×; synthetic MatMulNBits shapes up to 2.96× on M=128, N=2048, K=1536) while the symmetric control model is unchanged. The MLAS dispatch is factored cleanly around a single source-of-truth SelectKleidiAIQ4Backend(K, BlkLen, HasZp, config) selector, and the operator-side PrePack/shared-prepack machinery is extended to correctly decline packing when zero points arrive at runtime. Copilot's four review comments (two params{} value-initialization issues, two "scales with per-N temporaries" memory concerns) are all resolved in commit 60a6d9c — the author replaced full-N temporaries with a reused nr-wide panel scratch buffer, reporting a ~98% reduction in cumulative scratch allocation for the pack step under malloc_history -allEvents. CI on the latest force-push is still catching up (26/78 at time of writing; the prior head 724fae8 reached 59/86 without new failures), so pre-merge the CI-green gate should be confirmed on the current head.

What looks correct

1. Backend selection is single-source and consistent across call sites. SelectKleidiAIQ4Backend is queried by UsePacked_CompInt8, NeedsPackedZpCorrection_CompInt8, PackedQ4BitGemmNAlignment_CompInt8, QuantizeA_Packed_CompInt8, SQ4BitGemmKernel_Packed_CompInt8, QNBitGemmPackQuantBDataSize, SQ4BitGemmPackQuantBDataAndBlkSum, and QNBitGemmPerGemmWorkspaceSize. Every call site threads the same HasZp value from the caller's QuantBZeroPoint != nullptr — the "packed workspace is sized for X but decoded as Y" hazard that plagued the earlier symmetric-plus-correction shape doesn't recur, because a single boolean input governs the shape everywhere. IsKleidiAIQ4ShapeSupported (renamed from UseKleidiAI per the previous review round's cleanup) now gates only on shape (BlkLen > 0 && BlkLen % 32 == 0 && K % BlkLen == 0) and defers ISA gating to the per-backend arm in SelectKleidiAIQ4Backend. The asymmetric arm accepts SME2 || i8mm || dotprod; the symmetric arm accepts i8mm || dotprod (matching the existing symmetric kernel table, which never had SME2 variants). No regression to non-KleidiAI paths.
2. NeedsPackedZpCorrection_CompInt8 guards the correction step to exactly the symmetric backend. The old code applied C += AFloatBlkSum * BZpCorr^T whenever DataParams->BZpCorr != nullptr && DataParams->AFloatBlkSum != nullptr, which was only ever set by KleidiAI symmetric with zero points. The new RequiresPackedZpCorrection guard (which returns HasZp && backend == Qai8dxpQsi4c32p) makes this explicit and prevents a future backend from accidentally allocating the correction buffers and getting an unwanted mul-add. The Q4 asymmetric backend does NOT allocate BZpCorr/AFloatBlkSum (the workspace-size query in QNBitGemmPerGemmWorkspaceSize::Qsi8d32pQai4c32p does not append the M * BlockCountK float region), and the correction step is elided at compute time. Consistent.
3. Panel-based RHS packing is memory-safe and correct.
   • zero_offsets (size nr * BlockCountK floats) and rhs_for_kai (size nr * K/2 bytes) are allocated once outside the panel loop and reused. Peak scratch is bounded by kernel width (typically nr = 4) × K, not by N. On Qwen-2.5 1.5B and Phi-3.5 Mini this is single-digit KiB per node instead of many-MiB temporaries.
   • The final panel (panel_rows < nr) leaves stale rows in the reused buffers, but kai_run_rhs_pack_nxk_qai4c32p_* is invoked with n = panel_rows, so the KAI packer only reads the live rows. This is safe under the current KAI ABI; a defensive std::fill_n(zero_offsets.begin() + panel_rows * BlockCountK, ...) would harden against a future KAI change to read up to nr internally, but is not needed today.
   • The nibble swap ((rhs_panel[i] & 0x0F) << 4) | ((rhs_panel[i] & 0xF0) >> 4) correctly translates ORT's s1s0 (low nibble first) uint4 packing into KAI's s0s1 (high nibble first) qau4 layout. K is guaranteed even because BlkLen % 32 == 0 and K % BlkLen == 0, so K/2 is exact.
   • params.rhs_zero_point = 8 combined with kai_zp_offset = 8.0f - zp translates ORT's unsigned zp∈[0,15] into KAI's signed-int4 origin. zero_offsets[panel_idx] = kai_zp_offset * QuantBScaleBegin[src_idx] is the KAI-required "combined scale × (rhs_zero_point − zp)" preamble to the packer, which then internally folds it into the packed RHS. Byte-identical to reference.
4. StrideNThreadAlign integration. Threads must partition N on packed-panel boundaries so each thread reads its packed RHS from the correct n_step/nr-aligned offset. GetPackedQ4BitGemmNAlignment returns max(get_kai_n_alignment(gemm), get_kai_n_alignment(gemv)) where n_alignment = max(n_step, nr). StrideNThreadAlign = std::max(MLAS_QGEMM_STRIDEN_THREAD_ALIGN, PackedNAlignment) in MlasQNBitGemmBatch monotonically increases the pre-existing alignment (never decreases), so callers that used the default alignment are unaffected. Scoped to Variant == SQ4BitGemmVariant_CompInt8 — other variants keep their existing alignment. The assert(n_step != 0 && nr != 0 && (n_alignment % n_step) == 0 && (n_alignment % nr) == 0) catches a hypothetical future KAI kernel with mismatched n_step/nr — good defensive guard.
5. Shared packed B compatibility for the shared-prepack container. Constant zero points are folded into packed_b_ during PrePack(B) (via zp_ptr → MlasQNBitGemmPackQuantBData), so the framework's content-hash dedup sees byte-identical buffers across sessions, and the second session's UseSharedPrePackedBuffers correctly adopts the shared buffer without re-folding. The new TEST(MatMulNBits, SharedPrepackedWeights_AsymmetricPackedScales) exercises this positive path with 288 × 1024 × block_size=128 and asserts sharing works.
6. Runtime zero points correctly decline prepacking. RequiresDynamicZeroPointPrepackFallback(K, nbits, block_size, has_zp_arg, has_zp_input, compute_type, config) returns true when a zero_points input arg exists but is not a constant initializer AND the target shape would otherwise use the KleidiAI asymmetric packed path. Both MatMulNBits<T1>::PrePack (float) and MatMulNBits<MLFloat16>::PrePack short-circuit with return Status::OK() (no packed buffer, is_packed=false), so the framework falls through to ComputeBPacked's non-prepacked branch, which uses the reference path. UseSharedPrePackedBuffers also fails-fast with ORT_RETURN_IF(RequiresDynamicZeroPointPrepackFallback(...), "MatMulNBits cannot use shared prepacked B for KleidiAI Q4 with runtime zero_points. PrePack should have declined prepacking for this node.") so a shared B from a peer session that (somehow) got packed can't sneak in. Guarded by #if defined(MLAS_TARGET_ARM64) — non-ARM64 builds are compile-time-untouched.
7. UseSharedPrePackedBuffers now threads prepacked_buffer_sizes[0] into packed_b_size_. Previously the second session adopted the shared buffer without restoring its size, leaving packed_b_size_ at the default. Any downstream code that inspected packed_b_size_ (LUT-GEMM in particular) got a wrong answer. The PR passes the size through ORT_RETURN_IF(prepacked_buffer_sizes.empty(), "Missing MatMulNBits prepacked B buffer size metadata.") gated adoption. Latent bug fix.
8. MLFloat16 CompInt8 path is now unified with the fp32 path. The HQNBIT_CompInt8 compute-type is mapped to effective_compute_type = SQNBIT_CompInt8 for prepack and compute — MLFloat16 inputs are bulk-converted to fp32 at the operator boundary, and the same MLAS SQNBIT_CompInt8 pipeline (including KleidiAI asymmetric) is used. This eliminates the per-tile fp16→fp32 wrappers that used to sit inside HQNBIT_CompInt8 kernels and inherits KleidiAI support automatically for fp16 models. The packed_b_finalized_ bookkeeping is updated consistently so a shared-prepack second session in an fp16 model doesn't re-fold scales/ZP.
9. params{} value-initialization. The kai_rhs_pack_nxk_qsi4c32p_qsu4c32s1s0_params and kai_rhs_pack_nxk_qai4c32p_params structs are now {} value-initialized before their fields are set. This was Copilot's C1/C2 concern — the pre-fix code left params.scale_dt (in the asymmetric path) reading uninitialized memory, which the KAI ABI doesn't consult on that path but is still UB under the C++ standard. Zero-initialization removes the UB with no perf cost.
10. Test coverage is well-scoped.
    • TEST(MatMulNBits, SharedPrepackedWeights_AsymmetricPackedScales) — constant ZP + shared container → both sessions share, prepack folds scales/ZP into B.
    • TEST(MatMulNBits, DynamicZeroPoints_AsymmetricCompInt8) — runtime ZP, no shared container → PrePack declines, computation runs through the reference fallback with output_abs_error = 0.1f, output_rel_error = 0.02f.
    • TEST(MatMulNBits, SharedPrepackedWeights_DynamicZeroPoints_AsymmetricCompInt8) — runtime ZP + shared container → PrePack declines, kAddInitializerExpectNoPrepack verifier asserts number_of_pre_packed_weights_counter_session_1 == 0, number_of_elements_in_shared_container == 0, and the same for the second session's counters. Skipped cleanly (GTEST_SKIP()) when MlasQNBitGemmScalesPacked(1024, 4, 128, SQNBIT_CompInt8, true, nullptr) reports no active KleidiAI backend, so non-ARM64 and non-KleidiAI builds don't spuriously fail.
    • SQNBitGemmKleidiAIShortExecuteTest registers three MLAS unit tests: KleidiAIAsymGemv_M1_N257_K128 (partial trailing panel: N=257 ≠ multiple of 4), KleidiAIAsymGemm_M5_N257_K128 (partial panel + threadpool), KleidiAIAsymGemv_M1_N288_K1024_NoBias. #if defined(MLAS_TARGET_ARM64) gated with a GTEST_SKIP() fallback on other targets. Includes a MlasQNBitGemmScalesPacked(K, ...) availability probe before running so non-KleidiAI-builds skip cleanly.

Non-blocking observations

1. Verify has_zp_input_ naming semantics in the MatMulNBits constructor. The fallback correctness hinges on has_zp_input_ == false meaning "zero_points input is not a constant initializer" (i.e., runtime tensor) rather than merely "zero_points input is not connected". The DynamicZeroPoints_AsymmetricCompInt8 test with opts.zero_points_are_initializers = false implicitly relies on this: has_zp_arg_ = true (arg is connected) AND has_zp_input_ = false (not a constant) → fallback triggers. If the constructor sets has_zp_input_ = true whenever the input is connected regardless of const-ness, the fallback misses and the KleidiAI packer would run into assert(QuantBZPBegin != nullptr) (Debug) or a null deref (Release) via the subsequent TryGetConstantInput(...) → nullptr → zp_ptr = nullptr → MlasQNBitGemmPackQuantBData call. The test passing on zero_points_are_initializers = false is strong evidence the naming is "is-constant", but worth a one-line comment above the has_zp_input_ declaration to document this invariant for future contributors.

2. params.rhs_zero_point = 8 is a magic number. The KAI QAI4 signed-int4 origin (+8 offset that maps unsigned [0,15] into signed [−8,7]) is a KAI ABI constant. A named constexpr int kKaiQ4SignedZeroPointOffset = 8; at file scope (or as a static constexpr in namespace sqnbitgemm_neon) with a one-line comment would make the two arithmetic sites (params.rhs_zero_point = 8 and kai_zp_offset = 8.0f - zp) self-documenting and prevent them from drifting apart in a future refactor.

3. std::vector allocations in SQ4BitGemmPackQuantBDataAndBlkSum. zero_offsets (nr * BlockCountK floats) and rhs_for_kai (nr * K/2 bytes) are heap-allocated once per pack call. For a model with many MatMulNBits nodes this is a few KiB per node at load time — negligible in practice. Consider InlinedVector (per AGENTS.md convention) if peak load-time RSS ever becomes a concern, though the size is data-dependent so the small-buffer optimization won't always kick in.

4. RequiresDynamicZeroPointPrepackFallback has file-scope linkage. It's defined at file scope in matmul_nbits.cc without static or an anonymous namespace. If a future TU defines a function with the same signature, we'd get an ODR violation. Simple fix: move into the existing anonymous namespace at the top of the file, or mark static. Trivial.

5. switch (SelectKleidiAIQ4Backend(...)) pattern is repeated in ~5 places. Each dispatcher has the same three-arm switch — Qai8dxpQsi4c32p, Qsi8d32pQai4c32p, None. A small helper template <typename F> auto WithKleidiAIQ4Backend(K, BlkLen, HasZp, config, F&& sym, F&& asym, F&& fallback) that dispatches via a std::variant or if-else chain would collapse duplication, but the current form is more grep-friendly and reads directly. Deferrable style choice.

6. Symmetric backend on SME2-only hypothetical hardware. SelectKleidiAIQ4Backend with HasZp=false and SME2 alone (no dotprod, no i8mm) returns None, because has_neon_q4 = HasArmNeon_I8MM() || HasArmNeonDot(). No real SME2 chip lacks dotprod, so this is academic. If a symmetric SME2 KleidiAI kernel is ever added to the symmetric kernel table, this gating will need to be widened.

7. GetKleidiAIQ4LhsPackedOffset overload for the discarded size_t parameter. The symmetric ukernel offset getter is (m_idx, k); the asymmetric one is (m_idx, k, bl). The wrapper GetKleidiAIQ4LhsPackedOffset(const kai_matmul_clamp_f32_qai8dxp_qsi4c32p_ukernel& ukernel, size_t m_idx, size_t k, size_t) uses an unnamed size_t to keep a uniform three-arg call site. Clean, but a // bl unused for symmetric comment above the discarded param would help future readers.

8. KAI panel-based rhs_for_kai is fully overwritten per panel iteration. Confirmed by inspection: the loop for (size_t i = 0; i < rhs_panel_size; ++i) writes exactly panel_rows * rhs_stride bytes, and rhs_panel_size = panel_rows * rhs_stride. On the last (partial) panel, rhs_panel_size < nr * rhs_stride, so trailing nr - panel_rows rows in rhs_for_kai retain values from the previous iteration. As noted in point 3 of "What looks correct", the KAI packer reads only n = panel_rows rows, so this is safe under the current KAI ABI. A one-line // KAI reads only the first panel_rows rows; trailing rows in rhs_for_kai are not consulted. comment above the pack call would document the invariant.

9. The performance table format on the PR body would benefit from a "baseline path" column. The end-to-end 45.77 → 24.90 ms and synthetic 1.91×–2.96× uplifts are impressive but leave open which portion came from the SME2 kernels vs eliding the ZP correction step vs the shared-prepack path. A quick microbenchmark isolating each contributor would help reviewers gauge which environments benefit most (SME2 has the largest single win; the ZP-correction elision benefits all asymmetric configurations even on dotprod-only chips). Non-blocking for landing.

10. CI status. As of the latest force-push (60a6d9c), 26/78 checks are OK with the rest still in-flight. The prior head (724fae8) reached 59/86 without regressions. Recommend waiting for the current head's CI to complete a full pass before landing, particularly the ARM64 legs that exercise the KleidiAI packing paths.

Summary

Substantial and well-motivated integration of the KleidiAI asymmetric Q4 kernels. Correctness rests on a single SelectKleidiAIQ4Backend selector consistently propagated to every prepack/quantize/compute site, byte-identical shared-prepack outputs when zero points are constant, and a clean PrePack decline (with matching UseSharedPrePackedBuffers guard) when zero points are runtime tensors. Both the shape-driven and dynamic-ZP fallbacks have dedicated tests. Copilot's four review comments — two zero-initialization fixes and two per-N-temporary memory concerns — are addressed in commit 60a6d9c with the switch to nr-wide panel scratch buffers, and the author confirmed a ~98% reduction in cumulative scratch allocation. Two observations worth double-checking before merge: (a) that has_zp_input_ naming in the MatMulNBits constructor means "is-constant" and not merely "is-connected" (the DynamicZeroPoints_AsymmetricCompInt8 test is strong evidence it does), and (b) that CI on the current head reaches green across the ARM64 legs. Approve, no blockers.

[hariharans29]

Thanks you for this nice PR and for addressing the previous set of comments.

1. Can you please check if the failing CIs are related to your PR ?
2. I have one last set of review comments here- [MLAS] Integrate Q4 asymmetric KleidiAI kernels into MatMulNBits #29654 (comment). Please see if there is something low hanging that can be taken in with the intent of leaving the rest for a future PR ?