perf(ep-cpu): pre-pack QLinearMatMul's constant weight and translate signedness

[justinchuby]

What

QLinearMatMul re-packed the whole K x N weight inside MLAS on every call — mlas_qgemm_i32 left data.BIsPacked = false — and copied the weight into a fresh dense buffer first (to_dense_bytes(&inputs[3]), 12.8 MB at K=N=3584). At M=1 that setup was the call.

Separately, MLAS documents signed activations as unsupported off ARM, and u8 x i8 can saturate its i16 pair accumulator on AVX2 without VNNI, so those combinations were declined to a scalar loop and lost to ONNX Runtime by 5-6x.

Three changes:

1. Pre-pack a constant weight once. Bind MLAS's quantized MlasGemmPackBSize / MlasGemmPackB and a BIsPacked GEMM entry point. The kernel now takes set_constant_inputs and holds the pack for its lifetime, keyed on the weight's full identity — address, K, N, both signedness flags, and whether the bytes were sign translated. Once the pack exists the dense copy is skipped entirely.
2. Translate signedness instead of declining it. Any operand MLAS has no kernel for is moved into the unsigned domain: XOR 0x80 on its bytes, +128 on its zero point. The kernel computes sum_k (a_k - za)(b_k - zb), so shifting an operand and its zero point by the same constant leaves every i32 accumulator bit-identical, and the call lands on the u8 x u8 kernel this file already trusts as exact.
3. Requantize in parallel. Rows are independent, so above PARALLEL_MIN_WORK they run on par_chunks_mut into a pre-sized buffer instead of Vec::push per element, and a_scale * b_scale / y_scale is evaluated once per column rather than once per element when a_scale is per tensor (which the MLAS route already requires). The association is untouched, so the results are the same f32 bits.

Measurements

AMD EPYC 9V74 (32 vCPU / 16 physical cores, AVX2 + FMA + F16C, no AVX-512/VNNI/AMX), ONNX Runtime 1.27.0, same process, same binary, interleaved A/B, parity asserted on every rep, K = N = 3584, 3 warmups, 7 runs x 5 reps.

All ratios are ours / ORT — lower is better, < 1.00 means we are faster. Steady state (weight pre-packed on both sides).

Before this PR (main, 8 threads)

shape	u8 x u8 p50	i8 x i8 p50
M=1	20.59	5.20
M=128	4.02	4.97
M=512	2.28	5.22

(The i8 column is the declined scalar loop. An earlier round recorded 2.23-3.07 for that same unchanged path; the difference is the ORT 1.27.0 baseline and the harness on this shared host, not a regression — the kernel side did not change between the two rounds.)

After this PR

threads	u8 M=1 p50 / p90	u8 M=128 p50 / p90	u8 M=512 p50 / p90	i8 M=1 p50 / p90	i8 M=128 p50 / p90	i8 M=512 p50 / p90
1	1.13 / 1.18	1.20 / 1.21	1.20 / 1.21	0.03 / 0.03	0.25 / 0.25	0.26 / 0.26
4	1.80 / 2.21	2.18 / 2.53	2.07 / 2.12	0.05 / 0.08	0.39 / 0.44	0.38 / 0.40
8	2.33 / 2.77	2.43 / 2.80	2.12 / 2.22	0.09 / 0.09	0.47 / 0.53	0.42 / 0.43
16	2.34 / 2.58	2.65 / 3.42	2.08 / 2.18	0.10 / 0.10	0.60 / 0.69	0.42 / 0.47

Absolute milliseconds for the widest case, M=512, so the ratios can be divided out:

threads	ours u8	ORT u8	ours i8	ORT i8
1	121.28	100.45	124.13	469.78
8	33.63	15.99	34.57	82.57
16	26.82	13.08	27.47	65.82

i8 wins because ORT's own signed path is the slow one — ORT needs 469.78 ms where it needs 100.45 ms for the same shape in u8. After the translation both dtypes run the identical MLAS u8 x u8 kernel here at the same speed (124.13 vs 121.28 ms), so this is a real assignment win rather than an artefact of our two paths differing.

Cold vs steady

The benchmark harness runs a parity check before it starts timing, so it never sees a first call. Measured separately by qlinear_pack_cost_report (#[ignore]d), K = N = 3584, M = 1:

	time
first call, including the dense copy and the MLAS pack	6.35 ms
every later call	0.108 ms
never-packed path (what every call used to cost)	5.90 ms

The pack is repaid on the first call: +0.45 ms once, then -5.79 ms on every call after.

Memory: one MLAS packed panel per distinct (weight address, K, N, signedness), ~12.9 MB at K=N=3584. The kernel cache keys on node plus input shapes, so a model executed at several bucketed shapes holds one panel per bucket.

Assignment matrix change

dtype	region	before	after
i8 x i8	all	defer	claim (wins 1.7x-33x at every measured point)
u8 x u8	K*N >= 2^20	defer	defer (honest 1.13-2.65x, was 2.2-22x)
u8 x u8	K*N < 2^20	defer	claim (unmeasured; a partition boundary costs more than the ratio)
u8 x i8, i8 x u8	K*N >= 2^20	defer	defer (unmeasured; does not inherit the signed claim)
u8 x i8, i8 x u8	K*N < 2^20	defer	claim (unmeasured, like every other small shape here)

Both rules are measured on x86-64 AVX2 only and applied on every architecture, which is the convention the rest of that module uses. aarch64 has native i8 x i8 kernels (SDOT/SMMLA) that need no translation at all, so the claim there is if anything conservative — but its speed is unmeasured here and is not claimed to be measured. Correctness on that lane is covered unconditionally by qgemm_i32_matches_the_integer_oracle_for_every_signedness.

This does not make u8 fast. The residual u8 x u8 gap is thread scaling — 1.13-1.20x at one thread, widening to 2.08-2.65x at sixteen — which is the same root cause already documented for f32 dense and int4 MatMulNBits. It is deferred, and docs/performance/CPU_MATMUL_ASSIGNMENT.md says so.

Correctness

Nine new tests, and every new behaviour is falsified by a single injection (each was run with the guard removed and confirmed to fail):

injection	test that fails
flip_sign_domain uses 0x40 instead of 0x80	qlinear_matmul_reordered_accumulation_is_bit_identical
the zero point is not moved with the operand	qlinear_matmul_reordered_accumulation_is_bit_identical
the parallel requantize loses the row index	qlinear_matmul_reordered_accumulation_is_bit_identical
the pack key drops the weight address	a_different_weight_is_never_served_the_cached_pack
the pack ignores whether input 3 is constant	a_constant_weight_is_packed_once_and_reused
pack_key drops the batch guard	each_pack_guard_declines_on_its_own
pack_build drops the k * n length guard	each_pack_guard_declines_on_its_own
the weight is sized from input 1, not input 3	qlinear_matmul_defers_unsigned_activations_in_the_measured_region
the signed claim ignores the activation dtype	qlinear_matmul_claims_signed_activations
the signed pair is deferred as before	qlinear_matmul_claims_signed_activations

qlinear_matmul_claims_signed_activations also pins mixed signedness in both operand orders, above and below the measured region, so it cannot silently drift into either rule.

The bit-identity sweep now includes (96, 40, 900), which is above PARALLEL_MIN_WORK, so the parallel and serial row walks are both compared against a transcription of the previous loop, for both dtypes and per-tensor as well as per-axis quantization.

Three tests were also added in mlas-sys proving the packed quantized path is bit-identical to the unpacked one for every signedness and both zero-point layouts, with a checked > 0 guard so they cannot pass vacuously.

Validation

• cargo test -p onnx-runtime-ep-cpu --features mlas — 1269 pass
• cargo test -p onnx-runtime-ep-cpu (default features) — 1259 pass
• cargo test -p mlas-sys — 37 pass
• cargo clippy --all-targets clean with and without --features mlas (CI's Rust quality runs it without)
• cargo clippy --all-targets --target aarch64-unknown-linux-gnu clean
• cargo fmt --all -- --check clean
• NXRT_CPU_GEMM_BACKEND=generic (the x86-side proxy for the ARM lanes) leaves only the 3 pre-existing tests that assert auto-detection itself

Also here: the parallel-scaling gap is not the threadpool

The matrix has claimed since #1054 that the residual 1.4-2.4x loss at 2-16 threads on f32
dense and int4 MatMulNBits is "a threadpool/partitioning problem". Building the packed
QGEMM entry point above made that claim directly testable, and it is wrong.

Driving qgemm_i32_packed straight through this crate's work-stealing pool, K = N = 3584,
M = 512, warmed:

pool threads	1	2	4	8	16	32
MLAS GEMM only	99.9 ms	60.3 ms	26.9 ms	14.3 ms	13.6 ms	11.0 ms
speedup vs 1 thread	1.0x	1.7x	3.7x	7.0x	7.4x	9.1x

7.0x at eight threads and 7.4x at sixteen -- about what ORT achieves. mlas_threading_stats()
shows MLAS requesting exactly pool_threads partitions per call at every count, all of them
dispatched, and serial_fallback_calls == 0. The primitive and the pool both scale.

What does not scale is the serial per-call work around the primitive: densifying the
activation, allocating and zeroing the i32 accumulator, requantizing, and the executor's own
tensor handling. At M = 512 that is ~2 ms against a 100 ms GEMM at one thread (2%) and the
same ~2 ms against a 14 ms GEMM at eight (14%). Amdahl, not scheduling.

So the two commits after the review correct the attribution in
docs/performance/CPU_MATMUL_ASSIGNMENT.md and in assignment_policy.rs's
DECODE_PARALLEL_NOTE, and add qlinear_phase_report (#[ignore]d) which reproduces the
split. No deferral changes -- we still lose, so we still defer; only the stated cause, and
with it the direction of the open work, changes.

Honesty note: this host is shared and heavily contended. The GEMM and requantize columns above
are repeatable; the unattributed remainder varied 7.7 / 11.1 / 22.7 ms across three
consecutive runs of the identical configuration, so it is reported as a bound, not a
breakdown. Separately, the end-to-end harness reports 33.6 ms at M = 512, T = 8 where the
kernel alone measures 23.8 ms, implying ~10 ms/call of executor/session overhead outside the
kernel entirely. That is an unexplored lead, not a claim.

Not addressed here

• u8 x u8 thread scaling (deferred, documented).
• The ~10 ms/call executor overhead above.
• Real AVX-512 / VNNI / AMX and ARM64 silicon: unmeasurable on this host. The translation makes the fast path reachable on every ISA, and the exactness probe still gates u8 x i8 at runtime, so a VNNI or AMX host picks the native combination up automatically rather than through an ISA table.

[codecov]

Codecov Report

❌ Patch coverage is 85.24590% with 45 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.83%. Comparing base (7a29dc7) to head (43f5593).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
.../onnx-runtime-ep-cpu/src/kernels/qlinear_matmul.rs	58.25%	42 Missing and 1 partial ⚠️
crates/mlas-sys/src/lib.rs	98.52%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1086      +/-   ##
==========================================
- Coverage   80.39%   79.83%   -0.56%     
==========================================
  Files         369      369              
  Lines      161791   162043     +252     
  Branches   161791   162043     +252     
==========================================
- Hits       130072   129368     -704     
- Misses      26967    27926     +959     
+ Partials     4752     4749       -3     

Flag	Coverage Δ
cli-ort-linux	83.79% <ø> (ø)
cli-ort-windows	83.40% <ø> (+0.09%)	⬆️
mlas	84.13% <98.52%> (+0.89%)	⬆️
offline	79.62% <74.55%> (-0.61%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
...rates/onnx-runtime-ep-cpu/src/assignment_policy.rs	98.94% <100.00%> (+0.24%)	⬆️
crates/mlas-sys/src/lib.rs	83.17% <98.52%> (+0.93%)	⬆️
.../onnx-runtime-ep-cpu/src/kernels/qlinear_matmul.rs	84.16% <58.25%> (-1.61%)	⬇️

... and 7 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

⚠️ Benchmark Change Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
⚠️	add/large_f32_threads=1-internal/4194304	664.31 µs	769.15 µs	+15.8%
✅	add/medium_bf16_threads=1-internal/262144	95.50 µs	103.49 µs	+8.4%
✅	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	41.22 µs	43.90 µs	+6.5%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	8.71 ms	9.24 ms	+6.2%
✅	gather/medium_f32_threads=1-internal/32768	3.48 µs	3.66 µs	+5.4%
✅	add/small_f16_threads=1-internal/1024	418.0 ns	440.0 ns	+5.2%
✅	matmul/large_generic_bf16_threads=8/32x1024x1024	1.24 ms	1.29 ms	+4.0%
✅	add/small_f32_threads=1-internal/1024	186.8 ns	193.8 ns	+3.7%
✅	qwen3_sampling_processors/top_k_partial_selection	130.82 µs	135.61 µs	+3.7%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	354.43 µs	367.17 µs	+3.6%
✅	logit_processing/seven_processor_chain_per_step	300.58 µs	310.09 µs	+3.2%
✅	gather/small_bf16_threads=1-internal/4096	440.8 ns	454.5 ns	+3.1%
✅	add/medium_f16_threads=1-internal/262144	99.60 µs	102.60 µs	+3.0%
✅	kv_cache/alloc_dealloc_pages	36.72 µs	37.83 µs	+3.0%
✅	add/large_bf16_threads=1-internal/4194304	1.53 ms	1.58 ms	+3.0%
✅	matmul/small_generic_bf16_threads=8/1x256x256	29.72 µs	30.55 µs	+2.8%
✅	tokenization/decode_tokens_per_second	5.78 ms	5.92 ms	+2.5%
✅	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	517.49 µs	529.85 µs	+2.4%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.30 ms	5.42 ms	+2.2%
✅	add/large_f16_threads=1-internal/4194304	1.55 ms	1.58 ms	+2.2%
✅	matmul/medium_generic_f16_threads=8/32x512x512	28.37 µs	28.97 µs	+2.1%
✅	tokenization/encode_tokens_per_second	352.95 µs	360.15 µs	+2.0%
✅	reduce_mean/medium_f32_threads=1-internal/65536	228.35 µs	231.92 µs	+1.6%
✅	add/small_bf16_threads=1-internal/1024	416.4 ns	422.6 ns	+1.5%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.29 ms	3.34 ms	+1.3%
✅	qwen3_sampling_processors/top_k_top_p_fast	616.99 µs	624.90 µs	+1.3%
✅	add/medium_f32_threads=1-internal/262144	23.70 µs	23.93 µs	+1.0%
✅	matmul/small_generic_f16_threads=8/1x256x256	28.41 µs	28.66 µs	+0.9%
✅	reduce_mean/large_f32_threads=1-internal/262144	921.87 µs	929.33 µs	+0.8%
✅	gather/medium_f16_threads=1-internal/32768	2.25 µs	2.27 µs	+0.5%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.19 ms	2.20 ms	+0.5%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	162.08 µs	162.78 µs	+0.4%
✅	matmul/small_generic_f16_threads=1/1x256x256	28.47 µs	28.57 µs	+0.4%
✅	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	42.88 µs	42.98 µs	+0.2%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	496.00 µs	497.10 µs	+0.2%
✅	matmul/small_generic_bf16_threads=1/1x256x256	29.56 µs	29.60 µs	+0.1%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	1.99 ms	1.99 ms	+0.0%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.02 µs	14.02 µs	-0.0%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	1.86 ms	1.86 ms	-0.1%
✅	matmul/small_generic_f32_threads=1/1x256x256	35.07 µs	35.04 µs	-0.1%
✅	matmul/medium_generic_f16_threads=1/32x512x512	28.68 µs	28.64 µs	-0.2%
✅	gather/small_f16_threads=1-internal/4096	444.9 ns	443.9 ns	-0.2%
✅	grammar_masking/llguidance_compute_mask/32	71.51 µs	71.28 µs	-0.3%
✅	sampling_latency/greedy_per_token	3.06 µs	3.04 µs	-0.4%
✅	sampling_latency/top_p_per_token	360.19 µs	358.31 µs	-0.5%
✅	gather/large_f16_threads=1-internal/131072	11.16 µs	11.10 µs	-0.6%
✅	sampling_latency/top_k_per_token	49.74 µs	49.43 µs	-0.6%
✅	sampling_latency/min_p_per_token	196.72 µs	195.36 µs	-0.7%
✅	gather/medium_bf16_threads=1-internal/32768	2.26 µs	2.24 µs	-0.8%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	76.65 µs	75.93 µs	-0.9%
✅	matmul/small_generic_f32_threads=8/1x256x256	33.29 µs	32.82 µs	-1.4%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	499.56 µs	492.46 µs	-1.4%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	81.86 µs	80.46 µs	-1.7%
✅	gather/small_f32_threads=1-internal/4096	637.0 ns	623.1 ns	-2.2%
✅	matmul/medium_generic_f32_threads=8/32x512x512	928.45 µs	905.33 µs	-2.5%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	460.17 µs	445.75 µs	-3.1%
✅	gather/large_f32_threads=1-internal/131072	29.30 µs	27.08 µs	-7.6%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	4.05 ms	3.65 ms	-10.0%
✅	gather/large_bf16_threads=1-internal/131072	12.65 µs	10.78 µs	-14.7%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 3.07 3.03 3.12 }

What this cannot catch

• Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
• Sub-threshold regressions that compound over multiple PRs
• Performance changes that only manifest under GPU execution
• Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)

[justinchuby]

Three things in one PR, and they land in three different places under the strategy the owner set today. Reviewing them separately.

First, a blocker that applies regardless: the pre-pack is a fourth unaccounted resident weight buffer

"The kernel now takes set_constant_inputs and holds the pack for its lifetime" describes exactly the allocation class #1056 exists to govern: session-lifetime, scaled with weight size, invisible to the memory plan. The changed files are mlas-sys, assignment_policy.rs, qlinear_matmul.rs and a doc -- nothing in onnx-genai-engine, so the plan neither predicts nor admits these bytes.

We found three of these in the CPU path this week and they cost real money:

buffer	size	outcome
resident f32 dequant cache	~8x packed weight	took the 14B to ~66 GB before #979/#989
MLAS SQNBit packed buffer	~2x int4 bytes	peak RSS 8.17 → 25.5 GB on the 14B until #1051 governed it
weight transpose cache	1x f32 weight	governed by #1079 today

The rule from #1056: any allocation that outlives a single kernel call and scales with weight size must be declared to the plan before it is allocated, in the bytes actually allocated, and must be declinable. MlasGemmPackBSize already gives you the exact figure, and #1079 is a worked example of predictor + admission gate + decline path in this same crate. Please do the same here, including the two lessons that cost us cycles: account the bytes actually allocated (#1051's first attempt under-reported by 2.4x because it modelled them), and test against footprint, not against a proxy field staying empty.

Please also report peak RSS beside the timings. At K = N = 3584 the pack is not small, and a table of ratios with no memory column cannot show what it cost.

Second: change 2 is the best thing in this PR and it is ours, not MLAS's

The signedness translation is a genuine piece of knowledge, and the argument is airtight: the kernel computes sum_k (a_k - za)(b_k - zb), so shifting an operand and its zero point by the same constant leaves every i32 accumulator bit-identical. That turns a 5-6x loss into a supported path by arithmetic, not by dependency -- it would be just as true in a native kernel. Same for change 3's parallel requantize, which is our code already.

Third: change 1 deepens the MLAS dependency, and today's direction is the opposite

The owner's direction, set today: do not bundle MLAS by default; progressively absorb its optimizations into our native kernels. MLAS is a reference implementation we learn from, not a dependency we ship behind.

There is a hard fact behind that (filed as #1091): mlas is not in the default feature set of the CLI, the server, or onnx-runtime-ep-cpu, and the CPU kernels carry these routes behind #[cfg(feature = "mlas")] -- 24 gates in matmul.rs, 6 in qlinear_matmul.rs. So none of this PR's ours/ORT ratios describe what a user gets from a default build. That is not a criticism of your measurements, which are careful and honest about their configuration; it is that the configuration we measure is not the one we ship.

What I would like:

1. Say which native gap change 1 stands in for, and add it to the MLAS-routed speedups do not reach a default build, and the strategy is to absorb them natively #1091 absorption list. Binding MlasGemmPackB is a reasonable way to learn what the packed layout buys -- and quantifying that is genuinely useful -- but the destination is our own prepack, governed by the plan.
2. If you can, report the default-build number beside the MLAS one. "5-6x better with --features mlas" and "5-6x better" are different claims, and only one is about the product.
3. Keep changes 2 and 3 regardless of what happens to 1 -- they stand on their own and are not MLAS-dependent.

Not merging until the pre-pack is accounted and declinable. The rest I am happy with.