perf(cpu): symmetric int4 MatMulNBits takes the zero-copy borrowed decode path (#979)

[justinchuby]

Closes #979.

What

The CPU MatMulNBits decode dispatch has a zero-copy direct-int4 kernel
(borrowed_affine_int4_matmul) that reads packed weights in place. It was gated
on the presence of a zero_points input:

&& let Some(zero_points) = zero_points          // <-- fails for symmetric
&& let Some(zero_points) = contiguous_host_slice::<u8>(zero_points)

So symmetric int4 models (3 inputs, no zero_points) fell past it — and past
the accuracy_level == 4 packed path — into the m == 1 branch that builds and
retains weight_nk: OnceLock<Vec<f32>>, i.e. ~8x the file size in resident
RAM. The condition expressed the wrong intent; symmetric int4 is mathematically
simpler (its zero point is the implicit midpoint 8) yet got the expensive path.

The change

• Gate on symmetry explicitly, not on "a zero_points input happens to exist".
• borrowed_affine_int4_matmul and its #[cfg(target_arch = "aarch64")]
  NEON-dot sibling borrowed_affine_int4_matmul_m1_neon_dot now take
  Option<&[u8]> zero points and use the implicit midpoint when absent, via the
  existing NBitsLayout::zero_point (which already returns 1 << (bits-1) for
  None). The NEON correction (8 - zero_point) * activation_sum is 0 for
  symmetric — a no-op — so that fast path works unchanged for both cases.
• New borrow_optional_int4_zero_points helper: Some(None) for symmetric,
  Some(Some(zp)) for a borrowable asymmetric uint8 input, None only when a
  present zero_points tensor can't be borrowed in place (fall through, unchanged).

Asymmetric arithmetic is bit-identical (the only diff is Some(zp_row) →
zp_row). Symmetric adds an implicit zero point of 8, a mathematical identity.

Evidence

All CPU-only native build
(cargo build --release -p onnx-genai-cli --no-default-features --features native-backend),
--backend native, 64-token greedy decode, long prompt. Footprint =
PeakWorkingSet64 sampled to process exit on the same before/after harness
(contention-immune). "before" = branch base 532ef6bc; "after" = this branch.

1. Footprint (symmetric, qwen05b-q4, 3 inputs)

binary	peak working set
before	2,755 MB (matches the 2,753 MB in #979)
after	580 MB

For reference, the asymmetric qwen2.5-0.5b-q4_0-mobius on the same box measures
587 MB — i.e. symmetric now lands on the same borrowed path with the same
footprint as the model that already used it (580 vs 587 MB), down from 4.7x
worse. (My absolute after-number is higher than the 452 MB in #979 because the
native backend also stands up an ORT environment and the box is contended; the
contention-immune signal is the before→after drop on one harness.)

2. Numerics, symmetric — byte-identical before vs after. SHA-256 of the
64-token generation matched exactly (before == after == explicit-zero-points
oracle qwen05b-symzp).

3. Numerics, asymmetric — not regressed. qwen2.5-0.5b-q4_0-mobius:
byte-identical before vs after, footprint unchanged (587 MB → 587 MB).

4. Positive proof the path changed (not inferred from a footprint drop). New
unit test matmulnbits_symmetric_m1_borrows_instead_of_building_f32_cache:
executes a constant symmetric int4 kernel (the case that pre-#979 populated
weight_nk) and asserts weight_nk.get().is_none() and
!prepack_cache_populated(...), alongside a numeric match. The pre-existing
..._reuses_weight_for_new_activations test — which asserted symmetric constant
B builds and reuses a cache — is updated to assert it now borrows in place
(no cache), which is why its behavior flipped.

Gates (verbatim)

• cargo fmt -p onnx-runtime-ep-cpu — clean.
• cargo clippy -p onnx-runtime-ep-cpu -- -D warnings — exit 0, no warnings.
• cargo test -p onnx-runtime-ep-cpu:
  • test result: ok. 1066 passed; 0 failed; 10 ignored; 0 measured; 0 filtered out
  • test result: ok. 3 passed; 0 failed
  • test result: ok. 10 passed; 0 failed
  • test result: ok. 1 passed; 0 failed
  • test result: ok. 1 passed; 0 failed
  • test result: ok. 5 passed; 0 failed
  • test result: ok. 6 passed; 0 failed
  • test result: ok. 0 passed; 0 failed; 1 ignored
  • Total 1,092 passed; 0 failed (baseline 1,091 + 1 new test).
• cargo test -p onnx-genai-engine --lib:
  • test result: ok. 382 passed; 0 failed; 1 ignored; 0 measured; 0 filtered out

aarch64 / NEON note

The NEON-dot path (borrowed_affine_int4_matmul_m1_neon_dot) matters
disproportionately — the Windows-ARM/unified-memory users hit hardest by the 8x.
Its math is unchanged for symmetric (midpoint correction is 0). This x86 box
can't execute it, but I validated it compiles for the affected target:
cargo check -p onnx-runtime-ep-cpu --target aarch64-pc-windows-msvc → exit 0.
Runtime verification on aarch64 hardware is still pending.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 83.56%. Comparing base (2276c12) to head (e9b5ae3).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main     #989       +/-   ##
===========================================
+ Coverage   79.70%   83.56%    +3.86%     
===========================================
  Files         356       12      -344     
  Lines      150213     5270   -144943     
  Branches   150213     5270   -144943     
===========================================
- Hits       119721     4404   -115317     
+ Misses      25833      682    -25151     
+ Partials     4659      184     -4475     

Flag	Coverage Δ
cli-ort-linux	83.52% <ø> (ø)
cli-ort-windows	83.02% <ø> (-0.10%)	⬇️
mlas	?
offline	?

Flags with carried forward coverage won't be shown. Click here to find out more.
see 344 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[justinchuby]

Thanks — both measurements land, and they tighten this PR nicely. Responses in your order.

1. CUDA is untouched

Confirmed on my side too: this PR changes only crates/onnx-runtime-ep-cpu (git diff --stat is one file, kernels/matmul_nbits.rs). Nothing in onnx-runtime-ep-cuda moved. Your 1.59x GPU number is consistent with the CUDA MatMulNBits path never materialising symmetric weights, so there was nothing to fix there.

2. Two-backend oracle

Agreed, and I lean on exactly that: my after-binary generation for symmetric qwen05b-q4 is byte-identical (SHA-256) to the explicit-zero-points model qwen05b-symzp, which your GPU run independently confirms is a true identity rather than a CPU-kernel coincidence. If the two ever diverged, this implementation would be wrong.

Zero-byte advantage. Worth restating your point: the model-side workaround costs the added tensors (+7.7 MB, ~2%). This kernel fix needs no tensor at all — NBitsLayout::zero_point(None, _) returns the implicit midpoint 1 << (bits-1), so symmetric pays zero extra bytes and no model re-export.

3. Acceptance #4 — positive proof the branch is taken (leading with this)

You're right that a footprint drop alone is also consistent with a failed load or an unrelated path. I've made the proof positive and non-circular at two levels:

• Unit-level, deterministic (contention-immune): added #[cfg(test)] path counters split by symmetry, incremented at the borrowed branch. The regression test matmulnbits_symmetric_m1_borrows_instead_of_building_f32_cache now asserts the symmetric counter strictly increases across its single execute (positive: the borrowed zero-copy branch is the path taken), and weight_nk.get().is_none() and !prepack_cache_populated() (negative: no resident f32 / prepack cache ever built). The test uses constant B/scales — precisely the config that pre-Symmetric int4 has no zero-copy CPU kernel, so it pays 8x RAM: adding a constant zero_points=8 cuts footprint 6.1x with byte-identical output #979 populated weight_nk — so the negative assertions are non-vacuous. Counters are cfg(test) only, so the release binary is unaffected.
• Not a failed load: the after-binary produced a coherent, byte-identical 64-token generation on a long prompt. A model that failed to load cannot do that.

The test exercises MatMulNBitsKernel::execute — the same dispatch the real model calls — and qwen05b-q4's nodes share its attributes (bits=4, accuracy_level=0, symmetric, constant host initializers), so they route identically.

Conditions

All figures are PeakWorkingSet64 sampled to exit + SHA-256 byte-identity, both contention-immune, on one before/after harness (before = branch base 532ef6bc). I report no wall-clock. Gates re-run green after the test additions: ep-cpu 1,092 passed; 0 failed; engine --lib 382 passed; 0 failed; clippy -- -D warnings exit 0.

[github-actions]

🔴 Benchmark Regression 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
🔴	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	43.58 µs	127.83 µs	+193.3%
🔴	gather/large_bf16_threads=1-internal/131072	10.08 µs	23.48 µs	+133.0%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	509.20 µs	1.18 ms	+132.5%
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	43.91 µs	94.81 µs	+115.9%
🔴	matmul/large_generic_bf16_threads=8/32x1024x1024	1.28 ms	2.39 ms	+87.2%
🔴	gather/large_f16_threads=1-internal/131072	10.55 µs	18.75 µs	+77.6%
🔴	matmul/medium_generic_bf16_threads=8/32x512x512	388.28 µs	685.06 µs	+76.4%
🔴	gather/large_f32_threads=1-internal/131072	23.07 µs	37.64 µs	+63.2%
🔴	matmul/large_generic_f32_threads=8/32x1024x1024	4.27 ms	6.79 ms	+59.0%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	349.16 µs	546.46 µs	+56.5%
🔴	matmul/small_generic_bf16_threads=8/1x256x256	30.06 µs	42.85 µs	+42.6%
🔴	matmul/small_generic_bf16_threads=1/1x256x256	29.83 µs	39.03 µs	+30.8%
⚠️	matmul/small_generic_f16_threads=8/1x256x256	28.54 µs	36.53 µs	+28.0%
⚠️	matmul/medium_generic_f32_threads=1/32x512x512	2.19 ms	2.76 ms	+25.9%
⚠️	matmul/medium_generic_f32_threads=8/32x512x512	892.34 µs	1.09 ms	+22.3%
⚠️	matmul/large_generic_f16_threads=8/32x1024x1024	83.47 µs	102.04 µs	+22.3%
⚠️	matmul/medium_generic_f16_threads=8/32x512x512	27.58 µs	33.48 µs	+21.4%
⚠️	matmul/large_generic_f16_threads=1/32x1024x1024	77.17 µs	92.02 µs	+19.2%
⚠️	matmul/small_generic_f32_threads=1/1x256x256	39.76 µs	47.06 µs	+18.4%
⚠️	matmul/large_generic_f32_threads=1/32x1024x1024	8.66 ms	10.21 ms	+17.9%
⚠️	tokenization/encode_tokens_per_second	350.42 µs	410.80 µs	+17.2%
⚠️	matmul/medium_generic_bf16_threads=1/32x512x512	499.17 µs	585.07 µs	+17.2%
⚠️	matmul/medium_generic_f16_threads=1/32x512x512	28.03 µs	32.25 µs	+15.0%
✅	tokenization/decode_tokens_per_second	5.68 ms	6.43 ms	+13.1%
✅	add/medium_bf16_threads=1-internal/262144	2.63 ms	2.98 ms	+13.0%
✅	qwen3_sampling_processors/top_k_partial_selection	130.58 µs	147.51 µs	+13.0%
✅	qwen3_sampling_processors/top_k_top_p_fast	605.95 µs	681.73 µs	+12.5%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	1.95 ms	2.19 ms	+12.3%
✅	sampling_latency/top_k_per_token	48.84 µs	54.82 µs	+12.2%
✅	matmul/small_generic_f32_threads=8/1x256x256	35.91 µs	40.25 µs	+12.1%
✅	sampling_latency/top_p_per_token	360.45 µs	403.73 µs	+12.0%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	1.97 ms	2.17 ms	+10.3%
✅	sampling_latency/greedy_per_token	3.01 µs	3.32 µs	+10.3%
✅	sampling_latency/min_p_per_token	195.20 µs	214.97 µs	+10.1%
✅	logit_processing/seven_processor_chain_per_step	294.80 µs	322.40 µs	+9.4%
✅	kv_cache/alloc_dealloc_pages	36.42 µs	39.45 µs	+8.3%
✅	grammar_masking/llguidance_compute_mask/32	71.50 µs	77.21 µs	+8.0%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	74.00 µs	79.88 µs	+7.9%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	486.44 µs	524.76 µs	+7.9%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.31 ms	5.69 ms	+7.3%
✅	matmul/small_generic_f16_threads=1/1x256x256	30.84 µs	33.07 µs	+7.2%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.30 ms	3.53 ms	+7.0%
✅	add/large_f32_threads=1-internal/4194304	39.71 ms	40.60 ms	+2.2%
✅	reduce_mean/small_f32_threads=1-internal/4096	14.63 µs	14.95 µs	+2.2%
✅	gather/small_f32_threads=1-internal/4096	623.3 ns	630.8 ns	+1.2%
✅	gather/medium_f16_threads=1-internal/32768	2.29 µs	2.30 µs	+0.4%
✅	gather/medium_bf16_threads=1-internal/32768	2.28 µs	2.28 µs	-0.1%
✅	add/large_f16_threads=1-internal/4194304	41.86 ms	41.01 ms	-2.0%
✅	gather/medium_f32_threads=1-internal/32768	3.68 µs	3.58 µs	-2.6%
✅	gather/small_f16_threads=1-internal/4096	458.8 ns	440.9 ns	-3.9%
✅	gather/small_bf16_threads=1-internal/4096	486.5 ns	461.8 ns	-5.1%
✅	reduce_mean/large_f32_threads=1-internal/262144	966.10 µs	911.65 µs	-5.6%
✅	add/small_bf16_threads=1-internal/1024	13.11 µs	12.21 µs	-6.9%
✅	add/small_f16_threads=1-internal/1024	12.92 µs	12.02 µs	-7.0%
✅	add/medium_f16_threads=1-internal/262144	2.67 ms	2.43 ms	-9.2%
✅	add/large_bf16_threads=1-internal/4194304	48.55 ms	44.05 ms	-9.3%
✅	add/medium_f32_threads=1-internal/262144	2.62 ms	2.37 ms	-9.7%
✅	reduce_mean/medium_f32_threads=1-internal/65536	264.95 µs	226.62 µs	-14.5%
🟢	add/small_f32_threads=1-internal/1024	226.2 ns	180.4 ns	-20.2%

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.18 3.35 5.96 }

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)