[Copilot]

Description

• Save: Align every prepacked blob to 256 bytes, independently of optional page-alignment settings.
• Load: Retain aligned memory mappings; use aligned copies for legacy unaligned blobs or mapping failures.

Motivation and Context

CPU MatMulNBits with accuracy_level=4 can produce incorrect results on AVX2 after reloading saved prepacked weights. MLAS derives packed-region offsets from absolute pointer alignment, so reloading identical bytes at an unaligned address changes how the buffer is interpreted.

• Fixes [CPU EP] MatMulNBits (accuracy_level=4) produces wrong results on x86-64 AVX2 when loaded from session.save_external_prepacked_constant_initializers output (1.31.1) #33677

[azure-pipelines]

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

[Tianlei Wu (tianleiwu)]

Verdict: COMMENT — the fix is correct and verified end to end (the bug reproduces on main and is gone with this PR). The findings below are non-blocking hardening and maintainability suggestions.

What I verified (Linux x86-64, AMD EPYC 7V12 with AVX2, GCC 13.2, Release; built together with #33701, whose files are disjoint)

• PR head: MatMulNBits.*:MatMulNBitsLutGemm.* 47 passed / 6 skipped (the new end-to-end test ran and passed in all 4 configurations); TensorProtoUtilsTest.*:SessionStateTestSharedInitalizersWithPrePacking.*:SaveWithExternalInitializers.*:ExternalDataLoaderLifetimeTest.*:InferenceSessionTests.CompileApi* 56 passed.
• main production code + this PR's tests: MatMulNBits.SaveAndReloadPrepackedWeights fails in all 4 configurations ("65 of 65 differ"), and 3 TensorProtoUtilsTest tests fail (SetExternalDataInformation, PrepackedExternalDataHasDeterministicAlignedLayout, LoadPrepackedWeightsPreservesAlignment). The new tests do detect the bug.

#	Severity	Where	Issue
1	Suggestion	tensorprotoutils.cc (inline)	The copy fallback for misaligned mappings is silent and uncommented
2	Suggestion	tensor_external_data_info.cc (inline)	Name the MLAS absolute-alignment requirement as the reason for kAllocAlignment
3	Suggestion (optional)	MatMulNBits<T1>::UseSharedPrePackedBuffers in onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc (outside this diff)	The MLAS packed layout depends on the absolute address, and the failure mode is silently wrong numbers (65/65 outputs wrong in my repro). Consider a cheap check there that the incoming buffer is kAllocAlignment-aligned (return an error otherwise), or a follow-up that makes the packed layout base-relative, so any future producer of misaligned pre-packed buffers fails loudly
4	Nitpick	Tests	Legacy-file reload is covered only at unit level (offset 1), and the end-to-end test needs an int8 4-bit MLAS kernel (I could only exercise AVX2)
5	Nitpick	PR description	No validation evidence, and the checklist items are unchecked

Not run: Windows, macOS, ARM64, AVX-512/VNNI hardware, and the full onnxruntime_test_all / onnxruntime_provider_test suites.

[thiennq]

Reporter of #33677 here. I validated this PR on AVX-512, which the review above lists as not run.

Setup. Intel Xeon Gold 6133 (Skylake-SP; avx512f/bw/dq/cd/vl, no AVX512-VNNI / AVX-VNNI, so MLAS uses its MlasSQNBitGemmDispatchAvx512 path), Linux x86-64, CPUExecutionProvider only. The PR build is the onnxruntime-webgpu wheel from this PR's Linux WebGPU CI run, which reports git-commit-id=4f475ee (the CI merge of bc07b5d onto main). CI publishes no plain Linux CPU wheel for this PR, so I used this one. The wheel is tagged cp310 (CI builds it with CPython 3.10) but declares Requires-Python >=3.11, so I installed it on Python 3.10 with --ignore-requires-python; that mismatch is pre-existing and unrelated to this PR. The baseline is stock onnxruntime 1.31.1 on the same machine.

Repro sweep (repro_33677.py --sweep, int4, block 32). Each row saves the model with session.save_external_prepacked_constant_initializers, reloads it, and compares against an in-memory pack. consumed=yes means corrupting the blob bytes in the .data file changes the output, i.e. the disk blob is actually used.

case	1.31.1: blob offset % 64 / max_abs	this PR: blob offset % 64 / max_abs
M=1 K=64 N=33 acc=4	32 / 2.62	0 / 0
M=4 K=32 N=65 acc=4	16 / 5.58	0 / 0
M=4 K=32 N=67 acc=4	48 / 4.77	0 / 0
M=16 K=800 N=767 acc=4	48 / 6.09	0 / 0
control: M=4 K=64 N=33 acc=0	32 / 0	0 / 0
4 controls with aligned offsets (incl. the 2-layer transformer)	0 / 0	0 / 0

All 9 rows have consumed=yes with both builds. 1.31.1 exits with MISMATCH; this PR gives exact matches (max_abs 0) everywhere.

Unit tests from the same CI artifact, on the same machine (filter MatMulNBits.SaveAndReloadPrepackedWeights:TensorProtoUtilsTest.* on both binaries):

• onnxruntime_provider_test: MatMulNBits.SaveAndReloadPrepackedWeights 1/1 passed (ran, not skipped, so the CompInt8 kernel is available on this CPU).
• onnxruntime_test_all: TensorProtoUtilsTest.* 19/19 passed (including PrepackedExternalDataHasDeterministicAlignedLayout and LoadPrepackedWeightsPreservesAlignment).

Not covered here: AVX512-VNNI / AVX-VNNI kernels, Windows, macOS and ARM64.

[thiennq]

Copilot The AndroidBinarySizeCheckJob_MinimalBaseline check is failing on this PR, and it looks like a binary size threshold issue rather than a build or test failure:

RuntimeError: Sections total size for /build/1a/MinSizeRel/libonnxruntime.so of 1607802 exceeds threshold of 1607680 by 122. On-disk size=1586080

Failing job: https://github.com/microsoft/onnxruntime/actions/runs/40786732434/job/123169230991

The overage is only 122 bytes. It comes from the new alignment handling in GetFileContent (tensorprotoutils.cc), which adds an aligned-allocation fallback path that is compiled into the Android minimal build.

On this branch, .github/workflows/android.yml sets BINARY_SIZE_THRESHOLD_IN_BYTES=1607680. On main the same value is currently 1610752, which is 3072 bytes higher. With the main threshold this build (1607802 bytes) would pass with about 2950 bytes of headroom.

Please update BINARY_SIZE_THRESHOLD_IN_BYTES in .github/workflows/android.yml to 1610752 to match main (or merge/rebase the latest main into this branch so it picks up that value), and push the change so the check can re-run.

[Copilot]

🟢 Approval recommended

No active review issues block the changes preserving prepacked-weight alignment across model save and reload.

No open issues 🎉

What changed in this PR

Preserves prepacked-weight alignment across model save and reload, addressing the AVX2 MatMulNBits correctness issue reported in #33677.

Changes:

• Aligns every saved prepacked blob to 256 bytes, independently of optional page alignment.
• Retains aligned memory mappings and uses aligned copies for unaligned blobs or mapping failures.
• Adds alignment checks and MatMulNBits save/reload regression coverage.

File	Description
onnxruntime/​test/​framework/​tensorutils_test.cc	Verifies saved layout and loaded buffer alignment.
onnxruntime/​test/​contrib_ops/​matmul_4bits_test.cc	Checks MatMulNBits output equality after saving and reloading prepacked weights.
onnxruntime/​core/​framework/​tensorprotoutils.cc	Enforces load-time alignment with an aligned-copy fallback.
onnxruntime/​core/​framework/​tensor_external_data_info.cc	Aligns every prepacked blob before writing it.

_{Review effort: Balanced}

[Tianlei Wu (tianleiwu)]

Maintenance completion at 2517e4ca70c16f980442420a893e90b0eb6445be: feedback, alignment rationale/diagnostics, and main refresh are published; native provider/framework coverage passes, and all 101 exact-head checks are complete (100 success, one neutral).

The previously reported Android size failure is gone on the current head: AndroidBinarySizeCheckJob_MinimalBaseline succeeds. No test was skipped to hide that failure. Thank you for the additional AVX-512 verification.

There are no merge conflicts or unresolved review threads. Normal reviewer/branch-protection approval remains the final merge gate; this maintenance task did not merge the PR.