[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

Activity

1. Tianlei Wu (tianleiwu)

   assigned

   Copilot

   and

   Tianlei Wu (tianleiwu)

   on Oct 28, 2026
2. Copilot

   linked a pull request that will close this issuePreserve prepacked-weight alignment across save and reload #33700on Oct 28, 2026
3. thiennq commented on Nov 2, 2026

   

   thiennq

   on Nov 2, 2026

   Author

   More actions

   Follow-up from the reporter: I found a minimal, self-contained repro for this, along with the root cause.

   Summary. One MatMulNBits node (int4, block 32, accuracy_level=4, M=1, K=64, N=33) is enough to reproduce it; no fusion, M>1 or real model is needed. Prepacked blobs smaller than 1 MiB are written to the .data file without alignment. On x86-64, the accuracy_level=4 (CompInt8) packed layout locates its parts (packed B, block sums, scales) by rounding the buffer address up to 64 bytes. When memory-mapped on the tested x86-64 systems, the blob is used in place at an address with the same residue mod 64 as its file offset. For the affected accuracy_level=4 layout, a nonzero residue shifts the interpreted packed-data, block-sum and scale offsets. Moving each prepacked blob to a 64-byte aligned offset in the saved file removes the mismatch (see Realignment check).

   Results

   python repro_33677.py --sweep compares the reloaded prepacked model with the same model packed in memory. Both run the same kernel on the same packed bytes. Mismatch cells show max abs diff and cosine similarity. onnxruntime 1.31.1:

   Case	Blob offset % 64 (x86)	x86-64 AVX2	x86-64 AVX-512	arm64
   M=1, K=64, N=32, acc 4	0	match	match	match
   M=1, K=64, N=33, acc 4	32	1.74, cos 0.154	2.62, cos 0.187	match
   M=4, K=32, N=65, acc 4	16	5.58, cos 0.051	5.58, cos 0.051	match
   M=4, K=32, N=67, acc 4	48	4.77, cos -0.070	4.77, cos -0.070	match
   M=4, K=64, N=33, acc 0	32	match	match	match
   M=16, K=800, N=767, acc 4	48	7.22, cos 0.019	6.09, cos 0.019	match
   M=16, K=768, N=768, acc 4	0	match	match	match
   M=4, K=1056, N=1535, acc 4 (blob ≥ 1 MiB, aligned to 4096)	0	match	match	match¹
   2-layer transformer, M=16, K=256, N=1024, acc 4	0	match	match	match

   ¹ On arm64 the packed layout is smaller, so this blob is < 1 MiB, is written unaligned (offset % 64 = 48), and still matches.

   Platforms: x86-64 mismatches on a Windows VM (Intel/QEMU, AVX2), a Linux VM (AVX-512), and GitHub Actions ubuntu-latest (AMD EPYC 9V45/9V74, AVX-512) and windows-latest (EPYC 7763 AVX2 / 9V74 AVX-512), with both 1.31.1 and 1.32.0.dev20261031001. arm64 (macos-latest M2 Pro with both versions, and a local M1 with 1.31.1) matches in every case.

   The single-node probe in the original report may have passed because its blob happened to land on an aligned offset. In a real model, a blob whose length is not a multiple of 64 can leave subsequent small blobs misaligned; later data or padding can restore alignment.

   Repro

   pip install onnxruntime==1.31.1 onnx==1.23.2 numpy==2.5.4
   python repro_33677.py           # one node: K=64, N=33, M=1; exits 1 on mismatch
   python repro_33677.py --sweep   # matrix above, including the controls

   The script also checks that the blob on disk is actually used: ORT packs the weight again at load time and only swaps in the disk blob when the keys match, so a silent repack would also "match". The loader selects a disk blob by the key string stored in the proto and does not hash or checksum the blob bytes (the checksum field is written as 0 and not validated: tensor_external_data_info.cc L163, tensorprotoutils.cc L1741-L1762), so a corrupted blob is still selected. After the prepacked run the script therefore XORs every prepacked blob range in a copy of the .data file with 0xFF, reloads, and requires the output to change (consumed=yes). If a future version validated the bytes, the corrupted blob would be rejected and the case would be reported inconclusive, not as a pass. Every row above, including the controls, reports consumed=yes on the Linux AVX-512 VM and on the M1. Exit codes: 0 = all match and consumed, 1 = mismatch, 2 = no mismatch but a blob was not consumed (inconclusive).

   Repo with a CI workflow (ubuntu/windows/macos × 1.31.1/nightly): https://github.com/thiennq/ort-matmulnbits-prepack-repro

   repro_33677.py

   #!/usr/bin/env python3
   """Self-contained repro probe for onnxruntime issue #33677.
   
   Builds synthetic ONNX graphs with int4 MatMulNBits weights (random, seeded),
   saves them with ``session.save_external_prepacked_constant_initializers`` and
   compares the reloaded (prepacked) model against the same model loaded the
   normal way. Needs only numpy, onnx and onnxruntime.
   
   Finding: on x86-64 a MatMulNBits (bits=4, accuracy_level=4) weight computes
   wrong values after reload whenever its prepacked blob sits at a file offset
   that is not a multiple of 64. Blobs smaller than 1 MiB are written unaligned,
   right after the previous initializer. The MLAS x86 packed layout aligns its
   sub-buffers (packed B, block sums, scales) to 64 bytes relative to the
   buffer's own address. When memory-mapped on the tested x86-64 systems, the
   blob is used in place at an address with the same residue mod 64 as its file
   offset; for the affected accuracy_level=4 layout a nonzero residue shifts the
   interpreted packed-data, block-sum and scale offsets.
   One node and M=1 are enough; graph fusions play no part.
   
   Consumption check: ORT re-packs every weight at load time and only swaps in
   the blob from disk when its key (a hash of the freshly packed bytes) matches,
   so a silent in-memory repack would also "match". The loader selects a disk
   blob by the key string stored in the proto and does not hash or checksum the
   blob bytes (the checksum field is written as 0 and not validated), so a
   corrupted blob is still selected. After the prepacked run the script therefore
   XORs every prepacked blob range in a copy of the .data file with 0xFF, reloads,
   and requires the output to change (consumed=yes). If a future version validated
   the bytes, the corrupted blob would be rejected and the case would be reported
   inconclusive, not as a pass.
   
   Exit codes: 0 = every case matches and its blob was consumed, 1 = mismatch
   (bug reproduced), 2 = no mismatch but at least one case is inconclusive.
   
       python repro_33677.py                 # one MatMulNBits node, exit 1 on mismatch
       python repro_33677.py --sweep         # matrix incl. aligned / accuracy_level 0 controls
       python repro_33677.py --model transformer --m 16 --k 256 --n 1024
   """
   from __future__ import annotations
   
   import argparse
   import json
   import os
   import platform
   import shutil
   import sys
   import tempfile
   
   import numpy as np
   import onnx
   import onnxruntime as ort
   from onnx import TensorProto, helper, numpy_helper
   
   OPSET = [helper.make_opsetid("", 17), helper.make_opsetid("com.microsoft", 1)]
   
   
   # --------------------------------------------------------------------------- weights
   def quantize_int4(w: np.ndarray, block: int):
       """Symmetric block-wise int4 of w [K, N] in the MatMulNBits layout.
   
       Returns (packed uint8 [N, K/block, block/2], scales float32 [N, K/block],
       dequantized float32 [K, N]). No zero_points input: the kernel's default
       zero point for 4 bits is 8.
       """
       k, n = w.shape
       assert k % block == 0
       nb = k // block
       wt = w.T.reshape(n, nb, block)                       # [N, nb, block]
       scale = np.abs(wt).max(axis=2, keepdims=True) / 7.0
       scale[scale == 0] = 1.0
       q = np.clip(np.round(wt / scale), -8, 7).astype(np.int32) + 8   # 0..15
       packed = (q[..., 0::2] | (q[..., 1::2] << 4)).astype(np.uint8)  # low nibble first
       deq = ((q - 8) * scale).reshape(n, k).T.astype(np.float32)
       return packed, scale.reshape(n, nb).astype(np.float32), deq
   
   
   class Builder:
       def __init__(self, seed: int, block: int, acc: int, fp32: bool):
           self.rng = np.random.default_rng(seed)
           self.block, self.acc, self.fp32 = block, acc, fp32
           self.nodes, self.inits, self.i = [], [], 0
   
       def name(self, p: str) -> str:
           self.i += 1
           return f"{p}_{self.i}"
   
       def const(self, arr: np.ndarray, p: str = "c") -> str:
           n = self.name(p)
           self.inits.append(numpy_helper.from_array(arr, n))
           return n
   
       def node(self, op: str, inputs, domain: str = "", **attrs) -> str:
           out = self.name(op.lower())
           self.nodes.append(helper.make_node(op, list(inputs), [out], name=out + "_n",
                                              domain=domain, **attrs))
           return out
   
       def linear(self, x: str, k: int, n: int, bias: bool = True) -> str:
           w = (self.rng.standard_normal((k, n)) / np.sqrt(k)).astype(np.float32)
           packed, scales, deq = quantize_int4(w, self.block)
           if self.fp32:  # reference twin: same dequantized weights, plain MatMul
               y = self.node("MatMul", [x, self.const(deq, "w")])
           else:
               y = self.node("MatMulNBits", [x, self.const(packed, "qw"), self.const(scales, "s")],
                             domain="com.microsoft", K=k, N=n, bits=4, block_size=self.block,
                             accuracy_level=self.acc)
           if bias:
               y = self.node("Add", [y, self.const((0.02 * self.rng.standard_normal(n)).astype(np.float32), "b")])
           return y
   
       def layernorm(self, x: str, d: int) -> str:
           return self.node("LayerNormalization",
                            [x, self.const(np.ones(d, np.float32) + 0.05 * self.rng.standard_normal(d).astype(np.float32), "g"),
                             self.const((0.05 * self.rng.standard_normal(d)).astype(np.float32), "beta")],
                            axis=-1, epsilon=1e-5)
   
       def gelu(self, x: str) -> str:  # erf form, the pattern GeluFusion/BiasGelu match
           h = self.node("Div", [x, self.const(np.array(np.sqrt(2.0), np.float32))])
           h = self.node("Erf", [h])
           h = self.node("Add", [h, self.const(np.array(1.0, np.float32))])
           h = self.node("Mul", [x, h])
           return self.node("Mul", [h, self.const(np.array(0.5, np.float32))])
   
       def attention(self, x: str, d: int, heads: int) -> str:
           hd = d // heads
           shape = self.const(np.array([0, 0, heads, hd], np.int64), "shp")
   
           def split(t):
               return self.node("Transpose", [self.node("Reshape", [t, shape])], perm=[0, 2, 1, 3])
           q, k, v = (split(self.linear(x, d, d)) for _ in range(3))
           kt = self.node("Transpose", [k], perm=[0, 1, 3, 2])
           s = self.node("Mul", [self.node("MatMul", [q, kt]), self.const(np.array(1 / np.sqrt(hd), np.float32))])
           p = self.node("Softmax", [s], axis=-1)
           o = self.node("Transpose", [self.node("MatMul", [p, v])], perm=[0, 2, 1, 3])
           o = self.node("Reshape", [o, self.const(np.array([0, 0, d], np.int64), "shp")])
           return self.linear(o, d, d)
   
   
   def build(kind: str, args, fp32: bool) -> onnx.ModelProto:
       b = Builder(args.seed, args.block, args.acc, fp32)
       d, ff = args.k, args.n
       if kind == "single":  # one MatMulNBits node, 2-D input [M, K]
           x = helper.make_tensor_value_info("x", TensorProto.FLOAT, ["M", args.k])
           y = b.linear("x", args.k, args.n, bias=False)
       elif kind == "chain":  # MatMulNBits K->N then N->K, no other ops
           x = helper.make_tensor_value_info("x", TensorProto.FLOAT, ["M", args.k])
           y = b.linear(b.linear("x", d, ff, bias=False), ff, d, bias=False)
       elif kind in ("mlp", "transformer"):
           x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, "M", d])
           h = "x"
           for _ in range(args.layers):
               if kind == "transformer":
                   h = b.node("Add", [h, b.attention(b.layernorm(h, d), d, args.heads)])
               f = b.linear(b.gelu(b.linear(b.layernorm(h, d), d, ff)), ff, d)
               h = b.node("Add", [h, f])
           y = b.layernorm(h, d)
       else:
           raise SystemExit(f"unknown model {kind}")
       out_shape = [1, "M", d] if kind in ("mlp", "transformer") else ["M", args.n if kind == "single" else d]
       g = helper.make_graph(b.nodes, kind, [x], [helper.make_tensor_value_info(y, TensorProto.FLOAT, out_shape)],
                             b.inits)
       m = helper.make_model(g, opset_imports=OPSET, producer_name="ort-33677-repro")
       m.ir_version = 8
       return m
   
   
   # --------------------------------------------------------------------------- sessions
   def options(args, level=None) -> ort.SessionOptions:
       so = ort.SessionOptions()
       so.graph_optimization_level = level if level is not None else LEVELS[args.opt]
       so.intra_op_num_threads = args.threads
       so.log_severity_level = 3
       return so
   
   
   LEVELS = {"all": ort.GraphOptimizationLevel.ORT_ENABLE_ALL,
             "extended": ort.GraphOptimizationLevel.ORT_ENABLE_EXTENDED,
             "basic": ort.GraphOptimizationLevel.ORT_ENABLE_BASIC,
             "disable": ort.GraphOptimizationLevel.ORT_DISABLE_ALL}
   
   
   def save_prepacked(src: str, dst: str, args) -> None:
       """The flow from the issue: optimize + write packed weights to dst(.data)."""
       so = options(args)
       so.intra_op_num_threads = 1
       so.add_session_config_entry("session.save_external_prepacked_constant_initializers", "1")
       so.add_session_config_entry("session.optimized_model_external_initializers_file_name",
                                   os.path.basename(dst).replace(".onnx", ".data"))
       so.add_session_config_entry("session.optimized_model_external_initializers_min_size_in_bytes", "1024")
       so.optimized_model_filepath = dst
       ort.InferenceSession(src, so, providers=["CPUExecutionProvider"])
   
   
   def run(path: str, x: np.ndarray, args) -> np.ndarray:
       s = ort.InferenceSession(path, options(args), providers=["CPUExecutionProvider"])
       y = s.run(None, {s.get_inputs()[0].name: x})[0]
       del s  # release the session (and its mapping of the .data file) before temp files are deleted
       return y
   
   
   def compare(a: np.ndarray, b: np.ndarray) -> dict:
       a, b = a.ravel().astype(np.float64), b.ravel().astype(np.float64)
       if np.isnan(a).any() or np.isnan(b).any():
           return {"max_abs": float("nan"), "cos": float("nan")}
       cos = float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-30))
       return {"max_abs": float(np.abs(a - b).max()), "cos": cos}
   
   
   def prepacked_blobs(path: str) -> list[tuple[str, int, int]]:
       """(data file, offset, length) of every prepacked blob in the saved model."""
       blobs = []
       for t in onnx.load(path, load_external_data=False).graph.initializer:
           ext = {e.key: e.value for e in t.external_data}
           for key, value in ext.items():
               # "<key>|<offset>;<length>;<checksum>[|<offset>;<length>;<checksum>...]":
               # one key can hold several blobs. MatMulNBits writes one, but all are listed.
               if key.startswith("prepacked"):
                   for blob in value.split("|")[1:]:
                       off, length = blob.split(";")[:2]
                       blobs.append((ext["location"], int(off), int(length)))
       return blobs
   
   
   def run_corrupted(packed: str, blobs, x: np.ndarray, args) -> np.ndarray:
       """Run a copy of the prepacked model whose blob bytes are all XORed with 0xFF."""
       with tempfile.TemporaryDirectory() as tmp:
           dst = os.path.join(tmp, os.path.basename(packed))
           shutil.copyfile(packed, dst)
           src_dir = os.path.dirname(packed)
           for loc in {b[0] for b in blobs}:
               with open(os.path.join(src_dir, loc), "rb") as f:
                   data = bytearray(f.read())
               for _, off, length in (b for b in blobs if b[0] == loc):
                   data[off:off + length] = (np.frombuffer(data, np.uint8, length, off) ^ 0xFF).tobytes()
               with open(os.path.join(tmp, loc), "wb") as f:
                   f.write(data)
           return run(dst, x, args)
   
   
   def case(args) -> dict:
       with tempfile.TemporaryDirectory() as tmp:
           if args.keep:
               os.makedirs(args.keep, exist_ok=True)
               tmp = args.keep
           q, ref, packed = (os.path.join(tmp, f) for f in ("model.onnx", "ref.onnx", "model.packed.onnx"))
           onnx.save(build(args.model, args, fp32=False), q)
           onnx.save(build(args.model, args, fp32=True), ref)
           save_prepacked(q, packed, args)
           rng = np.random.default_rng(args.seed + 1)
           shape = (1, args.m, args.k) if args.model in ("mlp", "transformer") else (args.m, args.k)
           x = rng.standard_normal(shape).astype(np.float32)
           y_ref, y_q, y_p = run(ref, x, args), run(q, x, args), run(packed, x, args)
           blobs = prepacked_blobs(packed)
           # The disk blob is consumed only if corrupting it changes the output.
           consumed = bool(blobs) and not np.array_equal(run_corrupted(packed, blobs, x, args), y_p, equal_nan=True)
           r = {"model": args.model, "M": args.m, "K": args.k, "N": args.n, "block": args.block,
                "acc": args.acc, "layers": args.layers, "opt": args.opt, "threads": args.threads,
                "prepacked_offsets_mod64": [b[1] % 64 for b in blobs],
                "prepacked_lengths": [b[2] for b in blobs], "consumed": consumed,
                "unpacked_vs_fp32": compare(y_q, y_ref), "prepacked_vs_unpacked": compare(y_p, y_q)}
           # The prepacked weights run the same kernel on the same packed bytes, so
           # the two loads must agree (bit-exact on every platform where it works).
           d = r["prepacked_vs_unpacked"]["max_abs"]
           r["mismatch"] = not (d <= 1e-5)  # also true for NaN
           r["status"] = "MISMATCH" if r["mismatch"] else "ok" if consumed else "INCONCLUSIVE"
           return r
   
   
   def cpu_flags() -> str:
       try:
           with open("/proc/cpuinfo") as f:
               for line in f:
                   if line.startswith(("flags", "Features")):  # x86 / aarch64
                       fl = set(line.split(":")[1].split())
                       return " ".join(sorted(x for x in fl if x.startswith(("avx", "amx", "sve", "sme")) or
                                              x in ("f16c", "fma", "asimd", "asimddp", "asimdhp", "i8mm", "bf16")))
       except OSError:
           pass
       if sys.platform == "win32":
           try:
               import ctypes
               k = ctypes.windll.kernel32  # PF_AVX2_INSTRUCTIONS_AVAILABLE=40, PF_AVX512F=41
               return " ".join(n for n, f in (("avx", 39), ("avx2", 40), ("avx512f", 41))
                               if k.IsProcessorFeaturePresent(f))
           except Exception:
               return "?"
       if sys.platform == "darwin":
           import subprocess
           return subprocess.run(["sysctl", "-n", "machdep.cpu.brand_string"], capture_output=True, text=True).stdout.strip()
       return "?"
   
   
   # model:M:K:N:accuracy_level. The prepacked blob is written right after the
   # N*K/2-byte B initializer, so N*K/2 % 64 decides where it lands in the file.
   DEFAULT_SWEEP = " ".join([
       "single:1:64:32:4",     # B = 1024 bytes -> blob offset % 64 == 0  (control)
       "single:1:64:33:4",     # B = 1056 bytes -> blob offset % 64 == 32
       "single:4:32:65:4",     # B = 1040 bytes -> % 64 == 16
       "single:4:32:67:4",     # B = 1072 bytes -> % 64 == 48
       "single:4:64:33:0",     # same misaligned blob, accuracy_level 0 (control)
       "single:16:800:767:4",  # realistic size, % 64 == 48
       "single:16:768:768:4",  # realistic size, aligned (control)
       "single:4:1056:1535:4",  # B % 64 == 48 but blob >= 1 MiB, which the writer aligns to 4096 (control)
       "transformer:16:256:1024:4",  # 2-layer transformer, all blobs aligned (control)
   ])
   
   
   def main() -> int:
       ap = argparse.ArgumentParser()
       ap.add_argument("--model", default="single", choices=["single", "chain", "mlp", "transformer"])
       ap.add_argument("--m", type=int, default=1, help="rows of A (sequence length)")
       ap.add_argument("--k", type=int, default=64, help="K (model dim for mlp/transformer)")
       ap.add_argument("--n", type=int, default=33, help="N (FF dim for mlp/transformer)")
       ap.add_argument("--block", type=int, default=32)
       ap.add_argument("--acc", type=int, default=4)
       ap.add_argument("--layers", type=int, default=2)
       ap.add_argument("--heads", type=int, default=4)
       ap.add_argument("--opt", default="all", choices=list(LEVELS))
       ap.add_argument("--threads", type=int, default=1)
       ap.add_argument("--seed", type=int, default=0)
       ap.add_argument("--sweep", nargs="?", const="default",
                       help="run the default matrix, or a space-separated list of model:M:K:N:acc[:threads]")
       ap.add_argument("--json", action="store_true")
       ap.add_argument("--keep", help="write the models to this directory instead of a temp dir "
                                      "(one subdirectory per case with --sweep)")
       args = ap.parse_args()
   
       print(f"onnxruntime {ort.__version__}  onnx {onnx.__version__}  numpy {np.__version__}  "
             f"python {platform.python_version()}  {platform.system()} {platform.machine()}")
       print(f"cpu: {platform.processor() or '?'} | flags: {cpu_flags()}")
   
       if args.sweep:
           spec = DEFAULT_SWEEP if args.sweep == "default" else args.sweep
           grid = []
           for item in spec.split():
               f = item.split(":")
               g = dict(model=f[0], m=int(f[1]), k=int(f[2]), n=int(f[3]), acc=int(f[4]))
               if len(f) > 5:
                   g["threads"] = int(f[5])
               if args.keep:
                   g["keep"] = os.path.join(args.keep, item.replace(":", "_"))
               grid.append(g)
           results = [case(argparse.Namespace(**{**vars(args), **g})) for g in grid]
       else:
           results = [case(args)]
   
       def join(v):
           return ",".join(map(str, sorted(set(v)))) or "-"
   
       for r in results:
           if args.json:
               print(json.dumps(r))
               continue
           pv = r["prepacked_vs_unpacked"]
           print(f"{r['model']:<11} M={r['M']:<3} K={r['K']:<4} N={r['N']:<4} block={r['block']} acc={r['acc']} "
                 f"blob_off%64={join(r['prepacked_offsets_mod64']):<9} blob_len={join(r['prepacked_lengths']):<8} | "
                 f"prepacked vs unpacked: max_abs={pv['max_abs']:<9.3g} cos={pv['cos']:<9.5f} | "
                 f"consumed={'yes' if r['consumed'] else 'no':<3} | {r['status']}")
       if any(r["mismatch"] for r in results):
           print("RESULT: MISMATCH (bug reproduced)")
           return 1
       if any(r["status"] == "INCONCLUSIVE" for r in results):
           print("RESULT: INCONCLUSIVE (a prepacked blob on disk was not consumed; the match proves nothing)")
           return 2
       print("RESULT: all match (prepacked blobs consumed)")
       return 0
   
   
   if __name__ == "__main__":
       sys.exit(main())

   Root cause

   Permalinks are to v1.31.1. The same >= align_threshold condition is present on main at 2927f3b.

   1. Writer. ExternalDataInfo::WritePrepackedToFileAndAddToProto pads a prepacked blob only when align is set and the blob is at least align_threshold bytes; smaller blobs follow the previous data directly:
      

      onnxruntime/onnxruntime/core/framework/tensor_external_data_info.cc

      Lines 154 to 159 in 6af1a93

      	if (align && SafeInt<int64_t>(size_in_bytes) >= align_threshold) {
      	// return early on error
      	if (!AlignAndPad(os, on_disk_alignment, external_offset)) {
      	return os;
      	}
      	}

      
      Defaults are align_offset = false, align_threshold = 1048576, on_disk_alignment = 4096:
      

      onnxruntime/include/onnxruntime/core/graph/model_saving_options.h

      Lines 28 to 35 in 6af1a93

      	bool align_offset = false;
      	// Alignment threshold for size of data.
      	// Having a low threshold will waste file space for small initializers.
      	// Only when tensor's data size is > the page_align_threshold it will be force
      	// aligned. Default to 1MB.
      	int64_t align_threshold = 1048576;
      	// Alignment factor for big tensors (bigger than align_threshold). Defaults to 4K.
      	int64_t on_disk_alignment = 4096;

      
      The session save path sets align_offset = true and leaves the threshold at its default:
      

      onnxruntime/onnxruntime/core/session/inference_session.cc

      Lines 3258 to 3259 in 6af1a93

      	ModelSavingOptions model_saving_options{optimized_model_external_initializers_min_size_in_bytes};
      	model_saving_options.align_offset = true;

      
      Regular initializers get their natural element alignment, but nothing similar is applied before the prepacked blobs that follow them:
      

      onnxruntime/onnxruntime/core/graph/graph.cc

      Lines 5421 to 5427 in 6af1a93

      	// |<---previous data---->|<---padding--->|<----------------tensor data---------------->|
      	if (element_size > 1) {
      	ORT_RETURN_IF_NOT(ExternalDataInfo::AlignAndPad(external_stream, SafeInt<int64_t>(element_size),
      	external_offset),
      	"Failed writing natural alignment padding for external data to: ",
      	model_external_file_path);
      	}

      
      

      onnxruntime/onnxruntime/core/graph/graph.cc

      Lines 5461 to 5465 in 6af1a93

      	auto& os = ExternalDataInfo::WritePrepackedToFileAndAddToProto(
      	*prepacked_weights_for_graph_, blob_keys_to_external_data,
      	model_saving_options.align_offset, model_saving_options.align_threshold,
      	model_saving_options.on_disk_alignment,
      	external_stream, external_offset, *output_proto);

   2. Loader. LoadPrepackedWeightsFromFile calls GetFileContent, which memory-maps the range (or falls back to new char[] plus a read) and returns the pointer as is, without an alignment check. A mapping starts at a page (or allocation-granularity) boundary, so the mapped address keeps the file offset's residue mod 64:
      

      onnxruntime/onnxruntime/core/framework/tensorprotoutils.cc

      Lines 1683 to 1715 in 6af1a93

      	static Status GetFileContent(const RandomAccessFile& file, FileOffsetType offset,
      	size_t length, IAllocatorUniquePtr<void>& external_data) {
      	// query length if it is 0
      	if (length == 0) {
      	uint64_t file_length = 0;
      	ORT_RETURN_IF_ERROR(file.GetLength(file_length));
      	length = SafeInt<size_t>(file_length);
      	}
      	// first, try to map into memory
      	{
      	RandomAccessFile::MappedMemoryPtr mapped_memory{};
      	auto status = file.Map(offset, length, mapped_memory);
      	if (status.IsOK()) {
      	const auto mapped_memory_deleter = mapped_memory.get_deleter();
      	IAllocatorUniquePtr<void>::deleter_type raw_buffer_deleter =
      	[mapped_memory_deleter](void* p) {
      	mapped_memory_deleter(static_cast<char*>(p));
      	};
      	auto* mapped_data = mapped_memory.release();
      	IAllocatorUniquePtr<void> raw_buffer(mapped_data, std::move(raw_buffer_deleter));
      	external_data.swap(raw_buffer);
      	return Status::OK();
      	}
      	}
      	// if that fails, try to copy
      	auto buffer = std::make_unique<char[]>(length);
      	ORT_RETURN_IF_ERROR(file.Read(offset, gsl::make_span(buffer.get(), length)));
      	IAllocatorUniquePtr<void> raw_buffer(buffer.release(), [](void* p) { delete[] reinterpret_cast<char*>(p); });
      	external_data.swap(raw_buffer);
      	return Status::OK();

      
      

      onnxruntime/onnxruntime/core/framework/tensorprotoutils.cc

      Lines 1737 to 1764 in 6af1a93

      	static Status LoadPrepackedWeightsFromFile(const RandomAccessFile& file,
      	std::uintmax_t file_length,
      	const ExternalDataInfo::PrepackedInfos& prepacked_infos,
      	PrepackedWeightsForGraph& prepacked_info) {
      	for (const auto& [key, blobs] : prepacked_infos) {
      	PrePackedWeights prepacked_weights;
      	prepacked_weights.buffers_.reserve(blobs.size());
      	prepacked_weights.buffer_sizes_.reserve(blobs.size());
      	for (const auto& blob : blobs) {
      	const auto blob_offset = std::get<0>(blob);
      	const auto blob_length = std::get<1>(blob);
      	SafeInt<FileOffsetType> end_of_blob{blob_offset};
      	end_of_blob += blob_length;
      	ORT_RETURN_IF(blob_offset < 0 || static_cast<uintmax_t>(end_of_blob) > file_length,
      	"Pre-packed blob: ", key, " offset: ", blob_offset, " file_length: ", file_length,
      	" is out of bounds and can not read in full");
      	IAllocatorUniquePtr<void> data_ptr;
      	ORT_RETURN_IF_ERROR(GetFileContent(file, blob_offset, blob_length, data_ptr));
      	prepacked_weights.buffers_.push_back(std::move(data_ptr));
      	prepacked_weights.buffer_sizes_.push_back(blob_length);
      	}
      	if (!blobs.empty()) {
      	prepacked_info.InsertPrepackedWeights(key, std::move(prepacked_weights));
      	}
      	}
      	return Status::OK();

      
      MatMulNBits then uses the loaded buffer in place:
      

      onnxruntime/onnxruntime/contrib_ops/cpu/quantization/matmul_nbits.cc

      Line 894 in 6af1a93

      packed_b_ = std::move(prepacked_buffers[0]);

   3. MLAS layout. On AMD64, PackedQuantBDataStruct aligns PackedQuantBData to 64 by address, then aligns QuantBBlkSum to 64 by address, and the scales follow the block sums:
      

      onnxruntime/onnxruntime/core/mlas/lib/qnbitgemm.h

      Lines 75 to 77 in 6af1a93

      	#if defined(MLAS_TARGET_AMD64_IX86)
      	// avx512 requires alignment on a 64-byte boundary
      	PackedQuantBData = (std::byte*)MlasAlignAddress(PackedQuantBWorkspace, 64);

      
      

      onnxruntime/onnxruntime/core/mlas/lib/qnbitgemm.h

      Lines 95 to 104 in 6af1a93

      	QuantBBlkSum = (T*)(PackedQuantBData + PackedQuantBDataSize);
      	QuantBBlkSum = (T*)MlasAlignAddress(QuantBBlkSum, MlasQNBitQuantBBlkSumAlignment());
      	if (QuantAUnsigned) {
      	BlkUnsignedQuantAZeroPointCorrection = (T*)((std::byte*)QuantBBlkSum + BlkSumSize);
      	BlkUnsignedQuantAZeroPointCorrection = (T*)MlasAlignAddress(BlkUnsignedQuantAZeroPointCorrection, MlasQNBitQuantBBlkSumAlignment());
      	PackedQuantBScale = (T*)((std::byte*)BlkUnsignedQuantAZeroPointCorrection + BlkSumSize);
      	} else {
      	BlkUnsignedQuantAZeroPointCorrection = nullptr;
      	PackedQuantBScale = (T*)((std::byte*)QuantBBlkSum + BlkSumSize);

      
      At pack time the buffer comes from the CPU allocator, which aligns to at least MlasGetPreferredBufferAlignment() (64 here), so the initial base-address rounding is a no-op; internal section padding can still be needed:
      

      onnxruntime/onnxruntime/core/framework/allocator.cc

      Lines 170 to 172 in 6af1a93

      	void* CPUAllocator::Alloc(size_t size) {
      	const auto alignment = std::max(Info().device.GetAlignment(), MlasGetPreferredBufferAlignment());
      	return AllocatorDefaultAllocAligned(size, alignment);

      
      For K=64, N=33, block 32, the packed data is 1056 bytes and the block sums sit at 1088 relative to an aligned base. After reload at residue 32, the kernel reads packed data from offset 32 and block sums from offset 1120. This is wrong data, not just a slow unaligned load.

   This also explains the controls: accuracy_level=0 does not use the block-sum layout, blobs of 1 MiB or more are padded to 4096, and on arm64 the 4-bit accuracy_level=4 path never re-derives section pointers from the base address at compute time: the NEON dispatch has no SQ4BitGemmKernel_BlkSum_CompInt8 (qnbitgemm_kernel_neon.cpp L508 vs L512 for 8-bit), so MlasQNBitGemmBatch uses the packed pointer and the original scales directly:
   

   onnxruntime/onnxruntime/core/mlas/lib/qnbitgemm.cpp

   Lines 1604 to 1628 in 6af1a93

   	if (Variant == SQ4BitGemmVariant_CompInt8 && GetMlasPlatform().QNBitGemmDispatch->SQ4BitGemmKernel_BlkSum_CompInt8 != nullptr) {
   	PackedQuantBDataStruct<T, 4> packed_quant_b(const_cast<void*>(Data->QuantBDataWorkspace), N, BlockCountK, BlkLen, false);
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->PackedQuantBData = packed_quant_b.PackedQuantBData;
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->QuantBBlkSum = packed_quant_b.QuantBBlkSum;
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->QuantBScale = packed_quant_b.PackedQuantBScale;
   	PerGemmQuantAWorkspace per_gemm_quant_a_workspace(PerGemmWorkspace, M, BlockCountK, BlkLen);
   	ComputeOperation(BlkLen, K, Data, &per_gemm_quant_a_workspace, 0, M, 0, N, BackendKernelSelectorConfig);
   	} else if (Variant == SQ8BitGemmVariant_CompInt8 && GetMlasPlatform().QNBitGemmDispatch->SQ8BitGemmKernel_BlkSum_CompInt8 != nullptr) {
   	PackedQuantBDataStruct<T, 8> packed_quant_b(const_cast<void*>(Data->QuantBDataWorkspace), N, BlockCountK, BlkLen, GetMlasPlatform().ArmNeonIsQuantActivationsUnsigned);
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->PackedQuantBData = packed_quant_b.PackedQuantBData;
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->QuantBBlkSum = packed_quant_b.QuantBBlkSum;
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->QuantBScale = packed_quant_b.PackedQuantBScale;
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->BlkUnsignedQuantAZeroPointCorrection = packed_quant_b.BlkUnsignedQuantAZeroPointCorrection;
   	PerGemmQuantAWorkspace per_gemm_quant_a_workspace(PerGemmWorkspace, M, BlockCountK, BlkLen);
   	ComputeOperation(BlkLen, K, Data, &per_gemm_quant_a_workspace, 0, M, 0, N, BackendKernelSelectorConfig);
   	} else if (Variant == SQ2BitGemmVariant_CompInt8 && GetMlasPlatform().QNBitGemmDispatch->SQ2BitGemmKernel_BlkSum_CompInt8 != nullptr) {
   	PackedQuantBDataStruct<T, 2> packed_quant_b(const_cast<void*>(Data->QuantBDataWorkspace), N, BlockCountK, BlkLen, false);
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->PackedQuantBData = packed_quant_b.PackedQuantBData;
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->QuantBBlkSum = packed_quant_b.QuantBBlkSum;
   	const_cast<MLAS_QNBIT_GEMM_DATA_PARAMS<T>*>(Data)->QuantBScale = packed_quant_b.PackedQuantBScale;
   	PerGemmQuantAWorkspace per_gemm_quant_a_workspace(PerGemmWorkspace, M, BlockCountK, BlkLen);
   	ComputeOperation(BlkLen, K, Data, &per_gemm_quant_a_workspace, 0, M, 0, N, BackendKernelSelectorConfig);
   	} else {
   	ComputeOperation(BlkLen, K, Data, PerGemmWorkspace, 0, M, 0, N, BackendKernelSelectorConfig);

   arm64 is not necessarily safe in general, though: 8-bit CompInt8 on arm64 does use the BlkSum path, and the constructor aligns 8-bit packed data to 32 and the block sums to 64 by address, so bits=8 on arm64 is plausibly exposed to the same writer gap. I have not tested it here. This is another reason to fix it at the framework level rather than only for x86.

   Realignment check

   I rewrote the model.packed.data saved by --keep DIR so each prepacked blob starts at a 64-byte aligned offset (blobs ≥ 1 MiB kept 4096), with unchanged bytes; only the external-data file name and offsets in model.packed.onnx were updated (regular initializers were realigned too, which does not change their bytes). All three failing cases tested in this realignment experiment then had zero maximum absolute difference from the in-memory-packed model (Linux x86-64, AVX-512, onnxruntime 1.31.1):

   Case	Saved by ORT	Realigned (offset % 64 = 0)
   M=1, K=64, N=33	2.62	0
   M=16, K=800, N=767	6.09	0
   M=4, K=32, N=67	4.77	0

   In the sweep, every case with offset % 64 = 0 matches, accuracy_level=0 matches at residue 32, and arm64 matches at every offset represented in this sweep.

   Suggested fix

   1. Normalize at load time in or next to LoadPrepackedWeightsFromFile. If the pointer returned by GetFileContent (mmap or copy fallback) is not 64-byte aligned, copy the whole blob, same bytes and length, into a buffer from the CPU allocator. This fixes files already saved, and also the new char[] fallback, which does not guarantee 64-byte alignment whatever the file offset. Aligned mappings stay zero-copy. Check the returned pointer, not the file offset.
      A simpler alternative: ORT already runs PrePack in memory before it swaps in the disk blob (session_state.cc L702-L731), so a misaligned pointer could simply be rejected and the fresh in-memory pack kept. That has the same memory cost as the copy and needs less code.
      Either way, GetFileContent treats length == 0 as "map the whole file" (tensorprotoutils.cc L1686-L1690), so the loader fix needs an explicit guard for empty blobs.

   2. Writer: align every prepacked blob's starting offset to a multiple of 64, unconditionally, so new files load zero-copy. Don't gate this on align_offset, which ModelSavingOptions defaults to false. Keep the configured large-blob alignment. If it is validated as a positive power of two, the larger of it and 64 satisfies both requirements; otherwise validate the configuration or compute a checked common multiple. When the large-blob branch is inactive the configured alignment is unused, so callers that pass 0 there (as tensorutils_test.cc L234 does) keep working. The sizing and writing passes must keep producing identical layouts (see the comment at L142), so the padding belongs in this shared helper. Sketch assuming the active large-blob alignment has been validated as a positive power of two:

      // This x86 packed layout requires a 64-byte-aligned base.
      constexpr int64_t kMinPrepackedBlobAlignment = 64;
      
      for (size_t i = 0, size = prepacked_weights->buffers_.size(); i < size; ++i) {
        const auto size_in_bytes = prepacked_weights->buffer_sizes_[i];
        int64_t blob_alignment = kMinPrepackedBlobAlignment;
        if (align && SafeInt<int64_t>(size_in_bytes) >= align_threshold) {
          blob_alignment = std::max(blob_alignment, on_disk_alignment);  // requires validated positive power-of-two alignment
        }
        if (!AlignAndPad(os, blob_alignment, external_offset)) {
          return os;
        }
        // ... unchanged: record "offset;size;0", write the blob, advance external_offset
      }

   3. Tests:

      • Update PrepackedExternalDataHasDeterministicAlignedLayout in onnxruntime/test/framework/tensorutils_test.cc (L260), which expects the 3-byte blob to stay unaligned at offset 3. Also cover align = false and sizes below, at and above the threshold.
      • Add a save-and-reload numeric test: one MatMulNBits with K=64, N=33, block 32, accuracy_level=4, saved with session.save_external_prepacked_constant_initializers, reloaded in a new session and compared. Also check that the prepacked blob is actually consumed, so a silent repack cannot hide a regression: assert every prepacked offset recorded in the saved proto is a multiple of 64, and that number_of_prepacks / used_shared_buffers shows the disk blob was consumed.
      • Add a loader test with unchanged blob bytes at misaligned offsets (e.g. residues 1, 16, 32, 63) and a forced Map failure, asserting 64-byte aligned pointers and preserved bytes.

   Making the MLAS layout base-relative would be a larger change (SIMD alignment audit, compatibility with existing blobs), so I'd keep it separate.

   Workaround until a fix lands: regenerate the model and its prepacked artifact with accuracy_level=0 for MatMulNBits, or on x86 skip session.save_external_prepacked_constant_initializers and let ORT pack at load time.

   I'm happy to open a PR with (1), (2) and the tests if this direction looks right.
4. Tianlei Wu (tianleiwu)

   closed this as completedin #33700on Nov 10, 2026
5. Tianlei Wu (tianleiwu)

   added a commit that references this issue on Nov 10, 2026

   Preserve prepacked-weight alignment across save and reload (#33700)

   ...

   This issue will close once commit d0ce976 is merged into the 'main' branch.

   d0ce976