[MLAS] Reject Arm® KleidiAI™ Q4 prepack with dynamic scales - #32068
Open
Martin Klacer (martin-klacer-arm) wants to merge 2 commits into
Open
[MLAS] Reject Arm® KleidiAI™ Q4 prepack with dynamic scales#32068Martin Klacer (martin-klacer-arm) wants to merge 2 commits into
Martin Klacer (martin-klacer-arm) wants to merge 2 commits into
Conversation
… layer * Moved all KleidiAI-specific orchestration from qnbitgemm_kernel_neon into a new file `mlas/lib/kleidiai/qnbitgemm_kleidiai.cpp` * Changed the integration to go through an additional set of new overrides in platform.cpp, mirroring existing KleidiAI GEMM integrations * Removed obsolete zero-point correction in qnbitgemm_kleidiai Signed-off-by: Martin Klacer <martin.klacer@arm.com>
Signed-off-by: Martin Klacer <martin.klacer@arm.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Copilot started reviewing on behalf of
Hariharan Seshadri (hariharans29)
August 14, 2026 00:00
View session
Contributor
There was a problem hiding this comment.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.