[MLAS] Reject Arm® KleidiAI™ Q4 prepack with dynamic scales

[Martin Klacer (martin-klacer-arm)]

Description

The Arm® KleidiAI™ Q4 path requires B and its scales to be packed together during session initialisation. Since runtime-provided scales are unavailable at that point, attempting to prepack B can pass a null scale pointer to the RHS packer causing undefined behavior.

This change declines B prepacking for nodes with dynamic scales and ensures that the existing unpacked fallback is used at execution. This matches the existing handling of dynamic zero points.

Motivation and Context

This fixes a pre-existing bug uncovered during the Q4 refactor in #31720.

The PR is a stacked follow-up based on the branch for #31720 and should be merged after it. The changes on this branch relevant to this PR are in the commit 76d6137 that builds on top of the previous commit (b1c03d8) under review in #31720.

[azure-pipelines]

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

[Copilot]

Pull request overview

Rejects KleidiAI Q4 prepacking when scales are runtime-provided, ensuring execution uses the unpacked fallback.

Changes:

• Detects non-initializer scales and declines incompatible prepacking/shared buffers.
• Adds dynamic-scale fallback and sharing tests.
• Includes the stacked KleidiAI QNBitGemm backend refactor from #31720.

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc	Adds dynamic-scale prepack fallback.
onnxruntime/test/contrib_ops/matmul_4bits_test.cc	Tests runtime scales and sharing.
onnxruntime/test/mlas/unittest/test_sqnbitgemm.cpp	Expands KleidiAI GEMM coverage.
onnxruntime/core/mlas/lib/kleidiai/qnbitgemm_kleidiai.cpp	Implements the dedicated backend.
onnxruntime/core/mlas/lib/kleidiai/mlasi_kleidiai.h	Declares QNBit overrides.
onnxruntime/core/mlas/lib/kleidiai/kai_ukernel_interface.h	Defines kernel wrappers.
onnxruntime/core/mlas/lib/kleidiai/kai_ukernel_interface.cpp	Updates relocated interface includes.
onnxruntime/core/mlas/lib/qnbitgemm.cpp	Routes operations through overrides.
onnxruntime/core/mlas/lib/qnbitgemm.h	Removes obsolete packed dispatch APIs.
onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.cpp	Removes embedded KleidiAI orchestration.
onnxruntime/core/mlas/lib/qnbitgemm_kernel_neon.h	Removes relocated declarations.
onnxruntime/core/mlas/lib/sqnbitgemm_kernel_neon_int8.cpp	Removes relocated packed kernels.
onnxruntime/core/mlas/lib/platform.cpp	Registers KleidiAI overrides.
onnxruntime/core/mlas/lib/mlasi.h	Adds override function pointers.
onnxruntime/core/mlas/inc/mlas_qnbit.h	Removes obsolete correction fields.
cmake/onnxruntime_mlas.cmake	Builds the dedicated backend.

---

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

[Martin Klacer (martin-klacer-arm)]

Rebased on top of the latest refactor branch state (#31720), the first 3 commits are part of the refactor effort, the last commit (5587ec4) is the Q4 dynamic scales rejection change

[Hariharan Seshadri (hariharans29)]

Review — PR #32068: Reject KleidiAI Q4 prepack with dynamic scales

Scope clarification

Per author note, this PR is stacked on #31720 (the KleidiAI QNBitGemm backend refactor). Three of the four commits (d6dd652, 3c8b772, 5ee3070) belong to #31720 and appear here only because the branch was rebased. The actual change under review is the last commit, 5587ec4. I focused there.

The fix — bug and change

The KleidiAI Q4 PackKleidiAIQ4QuantBData (both symmetric Qai8dxpQsi4c32p and asymmetric Qsi8d32pQai4c32p paths) uses QuantBScale unconditionally: it reads through it, does bit-reinterpret to a uint16_t bf16 scale array, and passes it into the KleidiAI RHS packer as a required buffer. If B PrePack runs when scales are not an initializer, QuantBScale is nullptr at that point, and the pack path either dereferences null or does an in-place assert(QuantBScale != nullptr) — undefined behavior in a shipping build.

The change in matmul_nbits.cc generalizes the pre-existing "runtime-zero-points" prepack-decline into a runtime-quantization-parameter decline:

const bool has_runtime_quantization_parameter =
    !has_scales_initializer || (has_zp_arg && !has_zp_input);
return has_runtime_quantization_parameter && nbits == 4 && effective_compute_type == SQNBIT_CompInt8 &&
       MlasQNBitGemmScalesPacked(K, nbits, block_size, effective_compute_type,
                                 has_zp_arg, &backend_kernel_selector_config);

Correctness by case (I walked all four combinations):

scales initializer	has_zp_arg	has_zp_input	pre-PR result	post-PR result
✓	✗	—	prepack (symmetric OK)	prepack (unchanged) ✓
✓	✓	✓	prepack (asymmetric OK)	prepack (unchanged) ✓
✓	✓	✗	decline (runtime zp)	decline (unchanged) ✓
✗	✗	—	prepack (BUG: null scale)	decline ✓ (new)
✗	✓	✓	prepack (BUG: null scale)	decline ✓ (new)
✗	✓	✗	decline (runtime zp)	decline (unchanged) ✓

Behavior for the three unaffected rows is preserved; the two buggy rows now correctly fall back. ✓

The subtle-but-load-bearing MlasQNBitGemmScalesPacked argument change

The HasZeroPoint parameter passed into MlasQNBitGemmScalesPacked changes from unconditional true to has_zp_arg. This isn't a cosmetic touch — it's required by the extension:

• Pre-PR, the predicate could only be reached when has_zp_arg == true, so true was a safe (if slightly redundant) constant.
• Post-PR, !has_scales_initializer can be true with has_zp_arg == false (symmetric-with-runtime-scales). Querying MlasQNBitGemmScalesPacked(..., HasZeroPoint=true, ...) when the actual graph is symmetric would ask about the wrong KleidiAI backend (asymmetric Qsi8d32pQai4c32p) and could return the wrong answer for shapes where only the symmetric path is active.

The switch to has_zp_arg makes the query match the actual graph shape. Correctly paired with the predicate extension. ✓

Init and threading of has_scales_initializer_

Set in the constructor via info.TryGetConstantInput(InputIndex::scales, &tensor_scales). Exactly the same pattern as the existing has_zp_input_. Const-set-once bool. No lifetime or thread issue. ✓

Comment cleanup and error-message update

• Old comment references to BZpCorr are generalized to "zero-point-derived metadata" — accurate after the [MLAS] Refactor QNBit KleidiAI integration into dedicated layer #31720 refactor eliminates the standalone BZpCorr scratch.
• UseSharedPrePackedBuffers error message now says "runtime scales or zero_points" instead of just "runtime zero_points". Matches the widened predicate. ✓

Test coverage

New tests:

• DynamicScales_SymmetricCompInt8: runs a KleidiAI-symmetric-compatible shape (M=1, N=288, K=1024, block_size=128) with scales_are_initializers=false. Gated on MlasQNBitGemmScalesPacked(...) reporting the packed path is active. Verifies the unpacked path produces correct output. Fp32 and Fp16 both covered.
• SharedPrepackedWeights_DynamicScales_SymmetricCompInt8: uses PrepackSharingMode::kAddInitializerExpectNoPrepack combined with runtime scales. Verifies that a second session with the same B but runtime scales does NOT wrongly adopt a shared buffer packed with different scales.

Both are properly #if MLAS_TARGET_ARM64 guarded, with MlasQNBitGemmScalesPacked runtime shape-check inside so they auto-skip on hardware/build configs where KleidiAI isn't active.

The RunTest shared harness gets a new scales_are_initializers field to let existing sharing tests thread the initializer flag through consistently. Fine.

Missing coverage

Asymmetric with runtime scales is a real case that the widened predicate handles (has_scales_initializer=false && has_zp_arg=true → has_runtime_quantization_parameter=true), but no test explicitly exercises it. Adding a DynamicScales_AsymmetricCompInt8 (has_zero_point=true, zero_points_are_initializers=true, scales_are_initializers=false) would round out the matrix. Two-line change to a new TEST(...) block; worth adding before merge or as a fast follow-up.

Style nits

1. Function name: RequiresDynamicQuantizationParameterPrepackFallback reads mouthy. RequiresRuntimeQuantParamPrepackFallback or RequiresRuntimeParamsPrepackFallback scans better; the current name is already at the point where line-wrapping the call sites gets awkward. Non-blocking.

2. The comment // KleidiAI Q4 pack embeds scales and, for asymmetric weights, zero-point-derived metadata in B. — accurate; useful for the next reader. Good addition.

Non-issues I re-verified

• has_scales_initializer_ and has_zp_input_ both use the identical TryGetConstantInput pattern; semantics match. ✓
• The updated RequiresDynamicQuantizationParameterPrepackFallback continues to short-circuit on nbits == 4 && effective_compute_type == SQNBIT_CompInt8 before consulting MlasQNBitGemmScalesPacked, so no extra platform-detection cost on the common non-KleidiAI paths. ✓
• Test SharedPrepackedWeights_DynamicScales_SymmetricCompInt8 uses MakeSharingTestOptions(..., has_zero_point=false, ...) then flips scales_are_initializers=false. That combination is exactly the row-that-was-buggy. ✓
• The stacked-PR context means the buffer sharing content-hash comment updates ("BZpCorr" → "zero-point-derived metadata") only make sense after [MLAS] Refactor QNBit KleidiAI integration into dedicated layer #31720 lands. Since this PR merges after [MLAS] Refactor QNBit KleidiAI integration into dedicated layer #31720, the comments won't be stale at merge time.

Recommendation

Approve — pending an asymmetric-with-runtime-scales test to close the coverage matrix. The core change is small, precisely targeted, and correctly extends an existing defensive pattern to prevent a real null-pointer UB in the KleidiAI pack path. The has_zp_arg argument change to MlasQNBitGemmScalesPacked is subtle but load-bearing and correct.

Merge order note: this PR must merge after #31720. Author already flagged.