Govern the CPU MatMulNBits resident f32 decode cache (#971) by justinchuby · Pull Request #987 · justinchuby/onnx-genai

justinchuby · 2026-08-15T03:57:58Z

Summary

Closes #971.

The CPU MatMulNBits kernel keeps a resident dequantised f32 copy of the weights (weight_nk), built on the first m==1 decode step and held for the whole session. For models whose quantisation takes the generic f32 decode path this is ~8x the packed int4 weight. It buys ~2.4x decode throughput when it fits and collapses to paging (~50x slower) when it does not — yet it was taken unconditionally, and the memory plan reported the on-disk int4 size, so every fit/no-fit verdict and the derived KV budget were computed against a number ~8x too small.

This PR makes the trade truthfully accounted and governed, without removing the cache or changing what the kernels compute.

Part 1 — truthful accounting

The CPU EP (which owns the kernel dispatch) now exposes resident_dequant_f32_cache_bytes(&graph), backed by a predictor matmul_nbits_decode_caches_dequant_f32 that mirrors the kernel's branch order exactly.
The engine asks the EP rather than re-deriving the rule in the plan, so the accounting cannot silently drift from dispatch (the Non-CUDA platforms get a fabricated 8 GiB device capacity, so an 8.33 GB model reports fits=false on a unified-memory machine #947 failure mode). A control test (predictor_matches_actual_resident_cache) runs real kernels and asserts the predictor agrees with the weight actually cached — if dispatch changes and the predictor doesn't, the build fails.
MemoryStrategyPlan::total_weight_bytes now reflects the expansion the runtime will actually hold.

Part 2 — governed decision

build_memory_strategy_plan admits the cache only when the expanded footprint fits the residency budget; otherwise it declines, and the kernels dequantise on the fly (the pre-existing !can_prepack transient path — byte-identical math, nothing new written).
The choice and reason are reported as a resident_f32_weight_cache decision and in the memory strategy plan log line (resident_f32_cache_bytes, f32_weight_cache_admitted), mirroring the advisory_only visibility already on main.

Only the native CPU path passes non-zero expansion bytes. ORT and CUDA use different kernels, pass zero, always report admitted, and are unchanged. Explicitly not in scope: removing the f32 cache or making kernels consume quantised weights in place.

Acceptance evidence

Method: Process.PeakWorkingSet64 sampled to exit (monotonic; polled while alive, max kept), CPU-only release build (cargo build --release -p onnx-genai-cli --no-default-features --features native-backend), warm cache (first run discarded), --backend native --greedy, min peak of 2 repeats. Footprint leads; wall-clock is reported second because this shared box was under heavy contention during measurement (per the measurement-discipline skill).

arm	model	budget	`f32_weight_cache_admitted`	`total_weight_bytes` (plan)	peak WS
A	qwen05b-q4	none	true (admit)	2,334,951,891	2,752.3 MB
B	qwen05b-q4	`--vram-limit 1GiB`	false (decline)	359,107,027 (on-disk)	1,390.1 MB
C	qwen2.5-0.5b-q4_0-mobius	none	true (`resident_f32_cache_bytes=0`)	368,990,694	440.6 MB (1.19x)

Proof the paths engaged (native plan log line, verbatim):

Arm A (admit): total_weight_bytes=2334951891 resident_f32_cache_bytes=1975844864 f32_weight_cache_admitted=true
Arm B (decline): total_weight_bytes=359107027 resident_f32_cache_bytes=1975844864 f32_weight_cache_admitted=false resolved_device_budget_bytes=Some(1073741824)
Arm C (unaffected): total_weight_bytes=368990694 resident_f32_cache_bytes=0 f32_weight_cache_admitted=true

No budget / budget fits (Arm A): unchanged behaviour — peak 2,752.3 MB, exactly the 2,752 MB in the issue.
Budget below expanded footprint (Arm B): cache declined (proven by f32_weight_cache_admitted=false + the peak dropping from 2,752→1,390 MB), output byte-identical to Arm A ("Gravity is a fundamental force of nature that attracts objects with mass towards each other. It is the force").
Plan's reported weight bytes now reflect the expansion: before (= main's behaviour, still the declined/on-disk number) 359,107,027 → after (admitted) 2,334,951,891.
Symmetric model unaffected (Arm C): 440.6 MB ≈ 1.19x on-disk, resident_f32_cache_bytes=0 — it never took the f32 path.

Decode direction: under identical heavy contention, the declined arm (on-the-fly dequant) was dramatically slower per token than the admitted arm — the expected trade. Absolute ms/tok on this box is not meaningful today (contention inflated even the admitted arm well past the 105 ms/tok baseline), so per the measurement discipline only footprint is reported as authoritative.

Gates (verbatim)

cargo fmt -p onnx-genai-engine -p onnx-runtime-ep-cpu — clean, no changes.
cargo clippy -p onnx-genai-engine -p onnx-runtime-ep-cpu -- -D warnings — Finished, exit 0.
cargo test -p onnx-genai-engine --lib — test result: ok. 385 passed; 0 failed; 1 ignored (baseline 382 + 3 new governed-decision tests).
cargo test -p onnx-runtime-ep-cpu — all binaries ok; lib test result: ok. 1067 passed; 0 failed; 10 ignored; total 1093 passed; 0 failed across binaries (baseline 1091 + 2 new control tests).

codecov · 2026-08-15T04:06:47Z

Codecov Report

❌ Patch coverage is 69.44444% with 55 lines in your changes missing coverage. Please review.
✅ Project coverage is 78.98%. Comparing base (38f4581) to head (59d585c).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
...es/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs	69.44%	51 Missing and 4 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #987      +/-   ##
==========================================
- Coverage   79.62%   78.98%   -0.65%     
==========================================
  Files         356      356              
  Lines      150211   150390     +179     
  Branches   150211   150390     +179     
==========================================
- Hits       119608   118784     -824     
- Misses      25946    26949    +1003     
  Partials     4657     4657

Flag	Coverage Δ
cli-ort-linux	`83.52% <ø> (ø)`
cli-ort-windows	`83.02% <ø> (-0.10%)`	⬇️
mlas	`80.92% <ø> (-0.19%)`	⬇️
offline	`78.78% <69.44%> (-0.68%)`	⬇️

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

Files with missing lines	Coverage Δ
...es/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs	`74.16% <69.44%> (-6.75%)`	⬇️

... 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.

The CPU `MatMulNBits` kernel builds a resident dequantised f32 copy of the weights (`weight_nk`) on the first decode step and holds it for the whole session. For models whose quantisation takes the generic f32 decode path this is ~8x the packed weight. The trade buys ~2.4x decode throughput when it fits and collapses to paging (~50x slower) when it does not -- yet it was taken unconditionally, and the memory plan reported the on-disk int4 size, so every fit/no-fit verdict and KV budget was computed against a number ~8x too small. Part 1 -- truthful accounting. The CPU EP (which owns the kernel dispatch) now exposes `resident_dequant_f32_cache_bytes(&graph)`, backed by a predictor `matmul_nbits_decode_caches_dequant_f32` that mirrors the kernel's branch order exactly. The engine asks the EP rather than re-deriving the rule, so the accounting cannot silently drift from dispatch (the #947 failure mode). The discriminator is kernel dispatch (`bits`/`accuracy_level`/`group_indices`/`m`), not symmetry: a symmetric int4 node with no `accuracy_level` lands in the f32 cache, while an asymmetric zero-point int4 node takes the on-the-fly `borrowed_affine` path and caches nothing. A control test runs real kernels across that config matrix and asserts the predictor agrees with the weight actually cached, failing the build if dispatch changes without the predictor. `total_weight_bytes` now reflects the expansion the runtime will really hold. Part 2 -- governed decision. `build_memory_strategy_plan` admits the cache only when the expanded footprint fits the residency budget; otherwise it declines, and the kernels dequantise on the fly (the existing `!can_prepack` transient path, byte-identical math -- a lossless slower-per-token trade, never the lossy `accuracy_level=4` int8-activation path). The choice and its reason are reported as a `resident_f32_weight_cache` decision and in the `memory strategy plan` log line (`resident_f32_cache_bytes`, `f32_weight_cache_admitted`), mirroring the `advisory_only` visibility already on main. Only the native CPU path passes non-zero expansion bytes; ORT and CUDA use different kernels and pass zero, so they always report admitted and are unchanged. Models whose native CPU dispatch avoids the f32 cache (e.g. asymmetric zero-point int4) likewise report zero and are unaffected. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

justinchuby · 2026-08-15T04:11:25Z

Re: correction on the real discriminator (issuecomment-5299053308)

The correction doesn't change this PR's behaviour — the implementation was grounded on the actual zero_points/dispatch facts, not the sym/asym labels, so it already tracks the real discriminator. Reconciliation:

The discriminator is kernel dispatch, not symmetry — and the engine never re-derives it. resident_dequant_f32_cache_bytes calls the EP predictor matmul_nbits_decode_caches_dequant_f32, which mirrors the kernel's branch chain (bits/accuracy_level/group_indices/m). The control test predictor_matches_actual_resident_cache runs real kernels across accuracy_level ∈ {0,4} × asymmetric ∈ {false,true} and asserts the predictor equals actual weight_nk residency — so a dispatch change that isn't mirrored fails the build (the Non-CUDA platforms get a fabricated 8 GiB device capacity, so an 8.33 GB model reports fits=false on a unified-memory machine #947 guard you re-emphasised).
Verified end-to-end on the corrected facts. Symmetric qwen05b-q4 (no zero_points, no accuracy_level) → predicted+measured cache: resident_f32_cache_bytes=1975844864, peak 2752 MB. Asymmetric qwen2.5-0.5b-q4_0 (has zero_points) → borrowed_affine on-the-fly → predicted+measured no cache: resident_f32_cache_bytes=0, peak 440 MB.
The governed decline is LOSSLESS, never accuracy_level=4. Declining only toggles whether the transient dequant (dequantize_weight + gemv_nk) is retained; the math is identical (Arm B output was byte-identical to Arm A). No code path in this PR reads-or-sets accuracy_level as a memory fallback — the predictor only reads it to classify dispatch. I did not, and will not, reach for accuracy_level.
Fixed two doc/comment strings that carried the old (backwards) "symmetric == cheap" framing so they can't re-seed the drift.

One honest scope boundary (Part 1): resident_dequant_f32_cache_bytes accounts for the f32 weight_nk expansion only — the #971 subject. It returns 0 for the packed int4 path (accuracy_level=4), which as you measured still holds ~3.7x of the file (int8 repack etc.). No default model reaches that path (none carries accuracy_level), and that expansion is a different one that would need its own EP query. I left it out rather than silently guess at it; happy to scope it separately if you want the accounting to cover every path's real residency.

…ight-cache

github-actions · 2026-08-15T05:04:04Z

🔴 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
🔴	`qwen3_sampling_processors/top_k_full_sort_baseline`	2.12 ms	3.02 ms	+42.6%
🔴	`matmul/large_generic_f16_threads=1/32x1024x1024`	74.44 µs	101.18 µs	+35.9%
⚠️	`sampling_latency/top_p_per_token`	375.65 µs	449.53 µs	+19.7%
⚠️	`qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline`	3.27 ms	3.83 ms	+17.3%
✅	`qwen3_sampling_processors/top_k_top_p_full_sort_baseline`	5.22 ms	5.97 ms	+14.4%
✅	`qwen3_sampling_processors/top_k_partial_selection`	131.10 µs	148.60 µs	+13.3%
✅	`sampling_latency/min_p_per_token`	229.78 µs	256.25 µs	+11.5%
✅	`logit_processing/seven_processor_chain_per_step`	295.68 µs	327.89 µs	+10.9%
✅	`matmul/large_generic_f16_threads=8/32x1024x1024`	78.16 µs	86.63 µs	+10.8%
✅	`block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024`	40.32 µs	44.52 µs	+10.4%
✅	`block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1`	345.36 µs	381.11 µs	+10.4%
✅	`qwen3_sampling_processors/top_k_top_p_fast`	615.73 µs	673.59 µs	+9.4%
✅	`matmul/large_generic_bf16_threads=1/32x1024x1024`	1.80 ms	1.94 ms	+7.7%
✅	`grammar_masking/llguidance_compute_mask/32`	72.55 µs	77.57 µs	+6.9%
✅	`qwen3_sampling_processors/top_p_fast_after_top_k`	479.15 µs	508.09 µs	+6.0%
✅	`kv_cache/alloc_dealloc_pages`	36.69 µs	38.80 µs	+5.8%
✅	`matmul/medium_generic_bf16_threads=8/32x512x512`	352.37 µs	371.19 µs	+5.3%
✅	`sampling_latency/top_k_per_token`	51.53 µs	54.28 µs	+5.3%
✅	`matmul/medium_generic_bf16_threads=1/32x512x512`	480.91 µs	495.91 µs	+3.1%
✅	`sampling_latency/greedy_per_token`	2.96 µs	3.01 µs	+1.7%
✅	`matmul/medium_generic_f16_threads=8/32x512x512`	27.78 µs	28.24 µs	+1.6%
✅	`matmul/large_generic_f32_threads=1/32x1024x1024`	8.60 ms	8.63 ms	+0.4%
✅	`block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1`	72.55 µs	72.69 µs	+0.2%
✅	`tokenization/encode_tokens_per_second`	348.03 µs	347.41 µs	-0.2%
✅	`reduce_mean/small_f32_threads=1-internal/4096`	13.95 µs	13.83 µs	-0.9%
✅	`add/large_bf16_threads=1-internal/4194304`	39.61 ms	39.24 ms	-0.9%
✅	`block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024`	42.80 µs	42.39 µs	-1.0%
✅	`tokenization/decode_tokens_per_second`	5.73 ms	5.67 ms	-1.0%
✅	`add/large_f32_threads=1-internal/4194304`	37.84 ms	37.21 ms	-1.7%
✅	`add/large_f16_threads=1-internal/4194304`	38.60 ms	37.89 ms	-1.8%
✅	`matmul/large_generic_bf16_threads=8/32x1024x1024`	1.33 ms	1.30 ms	-1.9%
✅	`matmul/small_generic_f32_threads=8/1x256x256`	33.05 µs	32.15 µs	-2.7%
✅	`block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024`	538.97 µs	512.90 µs	-4.8%
✅	`matmul/large_generic_f32_threads=8/32x1024x1024`	3.85 ms	3.64 ms	-5.4%
✅	`add/small_bf16_threads=1-internal/1024`	12.82 µs	11.93 µs	-6.9%
✅	`add/medium_bf16_threads=1-internal/262144`	2.64 ms	2.45 ms	-7.2%
✅	`add/medium_f16_threads=1-internal/262144`	2.56 ms	2.38 ms	-7.3%
✅	`add/medium_f32_threads=1-internal/262144`	2.49 ms	2.31 ms	-7.4%
✅	`gather/small_f32_threads=1-internal/4096`	673.9 ns	622.8 ns	-7.6%
✅	`matmul/small_generic_bf16_threads=1/1x256x256`	31.96 µs	29.41 µs	-8.0%
✅	`matmul/medium_generic_f16_threads=1/32x512x512`	30.52 µs	27.83 µs	-8.8%
✅	`matmul/small_generic_f16_threads=1/1x256x256`	32.44 µs	28.77 µs	-11.3%
✅	`reduce_mean/medium_f32_threads=1-internal/65536`	255.45 µs	226.52 µs	-11.3%
✅	`matmul/medium_generic_f32_threads=8/32x512x512`	1.02 ms	897.66 µs	-11.8%
✅	`matmul/medium_generic_f32_threads=1/32x512x512`	2.42 ms	2.13 ms	-11.9%
🟢	`matmul/small_generic_f32_threads=1/1x256x256`	40.73 µs	33.86 µs	-16.9%
🟢	`gather/small_f16_threads=1-internal/4096`	531.4 ns	439.7 ns	-17.3%
🟢	`gather/large_f32_threads=1-internal/131072`	30.02 µs	24.74 µs	-17.6%
🟢	`matmul/small_generic_f16_threads=8/1x256x256`	35.43 µs	28.88 µs	-18.5%
🟢	`add/small_f16_threads=1-internal/1024`	14.68 µs	11.94 µs	-18.7%
🟢	`matmul/small_generic_bf16_threads=8/1x256x256`	35.88 µs	28.90 µs	-19.5%
🟢	`reduce_mean/large_f32_threads=1-internal/262144`	1.15 ms	910.71 µs	-20.5%
🟢	`gather/medium_bf16_threads=1-internal/32768`	2.87 µs	2.19 µs	-23.7%
🟢	`gather/small_bf16_threads=1-internal/4096`	648.9 ns	444.3 ns	-31.5%
🟢	`add/small_f32_threads=1-internal/1024`	273.7 ns	183.1 ns	-33.1%
🟢	`gather/large_bf16_threads=1-internal/131072`	15.83 µs	9.82 µs	-37.9%
🟢	`gather/medium_f32_threads=1-internal/32768`	5.63 µs	3.36 µs	-40.2%
🟢	`gather/large_f16_threads=1-internal/131072`	15.64 µs	9.15 µs	-41.5%
🟢	`gather/medium_f16_threads=1-internal/32768`	4.10 µs	2.19 µs	-46.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: { 4.02 3.27 5.38 }

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)

Adds the day's milestones: fabricated device/host capacity removed (#947/#952), symmetric int4's missing zero-copy CPU kernel and its 8x resident footprint (#979, 14B 66 GB -> 7.80 GB), the memory plan accounting for the file rather than the residency and the cache becoming a governed decision (#971/#987), advisory_only made visible (#955), native CPU load found serial rather than expensive (#959/#970, 6.4 s -> 3.5 s), ragged batch geometry (#750 stage 3a), the plugin EP claiming 0 of 585 nodes because kernels were advertised in the wrong domain (#956), two ruled-out hypotheses for the interspersed-partition hang (#982), and placement plus the cost model that would decide it (#994/#995) with the first measured rate (host DRAM 49.28 GB/s at 20 threads). Also records the method note, because it was the day's actual lesson: several of my own intermediate conclusions were wrong and retracted -- a validation that was clean because the knob never engaged, a load/decode split contaminated by a cold cache, a workaround inferred from directory names that was exactly backwards, and a footprint lever that looked free on a short prompt and diverged on a long one. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624

#987 added `matmul_nbits_decode_caches_dequant_f32` so the engine can ask the EP whether a model will build the resident f32 weight cache, and size its budget against the real footprint. #979 then extended the zero-copy borrowed int4 path to symmetric weights. Both are correct; merged together the predictor started lying. Branch (A) still required `has_zero_points`, which was the pre-#979 gate. The kernel now takes the borrowed path for symmetric and asymmetric alike (`borrow_optional_int4_zero_points` yields `Some(None)` for symmetric, meaning the implicit midpoint 8). So for every symmetric int4 model the predictor claimed a resident f32 cache that is no longer built. That is not only a red test. The governed cache decision (#987) sizes `resident_f32_cache_bytes` from this predictor, so it would over-report the footprint of symmetric models -- exactly the models #979 just made cheap -- and could decline a cache that was never going to exist. Drop the symmetry requirement so the predictor matches the kernel gate, and say in the comment why symmetry must not appear there, so the next change to either side has the coupling in front of it. `predictor_matches_actual_resident_cache` now passes. Its sibling `declining_resident_cache_is_byte_identical_and_holds_no_expansion` still fails: its premise ("bits=4, accuracy_level=0, symmetric -> the generic weight_nk path") is what #979 invalidated, so it needs a config that still caches. Both tests were already failing on main before this commit; this fixes one and narrows the other to a stated cause. The control test doing its job is the good news here: it was written to fail the build if the predictor drifted from the kernel, and that is precisely what it did. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624

…nk (#1008) ## Repair the #987/#979 interaction: symmetric int4 no longer caches Follow-up to #971 / #987. PR #979 extended the zero-copy borrowed int4 path to **symmetric** weights (`borrow_optional_int4_zero_points` → `Some(None)` for the implicit midpoint-8 case), so symmetric int4 with no `g_idx` no longer builds the resident `weight_nk` f32 cache. That invalidated the *premise* of one control test from #987. The predictor drift itself was already fixed on `main` (`2bf41eed`, symmetry removed from branch A). The remaining red test was `kernels::matmul_nbits::tests::declining_resident_cache_is_byte_identical_and_holds_no_expansion`, whose premise asserted that bits=4/accuracy_level=0/symmetric *caches* — no longer true. ### What changed - **Repaired the failing test's premise (did not delete it).** Switched the probe to a configuration that genuinely still reaches the `weight_nk` cache after #979: a constant `g_idx` present at input 4. Every earlier branch in `execute` (borrowed int4, 2-bit, `accuracy_level==4`, 8-bit) is gated on `group_indices.is_none()`, so a present `g_idx` falls through to the `m == 1` `weight_nk` cache. The g_idx values are the natural per-block identity (`depth / block_size`, in range `0..k_blocks`), so the decode math is unchanged and the byte-identity assertion between the cached and declined arms still holds. Symmetric configs carry no `zero_points`, so input 3 is an explicit `TensorView::absent(DataType::Uint8)` placeholder with `g_idx` at input 4. - **Strengthened the anti-drift control** `predictor_matches_actual_resident_cache`. After #979, every config the sweep previously exercised (`accuracy_level ∈ {0,4}` × `asymmetric ∈ {false,true}`, no g_idx) now *declines* the cache, so nothing validated the predictor's positive arm. Added a `has_g_idx ∈ {false,true}` dimension so at least one config family actually caches and the kernel-vs-predictor agreement is checked on both sides of the branch — restoring the positive-cache guard. ### Audit for other stale "symmetric int4 caches" assumptions Checked the rest of #987: - The predictor's own doc comment was already corrected on `main` (branch A, "Symmetry is deliberately NOT part of this condition… #979"). - `crates/onnx-genai-engine/src/engine/memory_strategy.rs`: the `resident_f32_cache_bytes` doc already states the discriminator is kernel dispatch "*not* symmetry" and its example (asymmetric int4 taking the borrowed path) remains accurate. No change needed. - No other comment, doc string, or accounting path asserted that symmetric int4 caches. ### Gates - `cargo fmt -p onnx-runtime-ep-cpu` — clean - `cargo clippy -p onnx-runtime-ep-cpu -- -D warnings` — 0 warnings - `cargo test -p onnx-runtime-ep-cpu --lib` — `test result: ok. 1068 passed; 0 failed; 10 ignored; 0 measured; 0 filtered out` - `cargo test -p onnx-genai-engine --lib` — `test result: ok. 385 passed; 0 failed; 1 ignored; 0 measured; 0 filtered out` Relates to #971. Repairs the #987/#979 interaction. Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com>

…fer (#1027 follow-up) (#1051) ## What Follow-up to **#1027** (int4 `accuracy_level=0` → MLAS SQNBit CompFp32 dispatch). #1027 shipped the speedup but left the MLAS packed weight buffer (~2x the int4 bytes, held beside the still-resident mmap) **invisible** to the memory ledger — the engine printed `resident_f32_cache_bytes=0` in both the borrowed and the MLAS route, so the governor could neither see nor bound the ~2x weight footprint. This PR implements the coordinator's requested changes on top of the merged dispatch. > **Note:** #1027 was squash-merged into `main` (commit `410bdea0`) while this revision was building, so the changes could not be pushed as a revision of #1027. This branch is cut from current `main` (which already contains the #1027 dispatch) and carries only the accounting/governance/test changes. ## Changes 1. **Account.** `matmul_nbits_resident_side_cache_bytes` returns, per constant-weight node, the MLAS `sqnbit_packed_b_size` when it takes the acc0 route, else the f32 dequant cache (`N*K*4`, #971), else 0. `resident_dequant_f32_cache_bytes` sums it over the graph, so the same `resident_f32_cache_bytes` the plan already consumes now includes the packed buffer. Accounting gates only on the route being reachable (env + a real MLAS kernel), never on the admission flag — the cost cannot vanish exactly when the plan needs it. 2. **Govern.** The plan's existing admission verdict (`f32_weight_cache_admitted`) now also drives `set_mlas_sqnbit_packing_enabled`. When declined (expanded footprint over budget), the route yields to the pre-existing borrowed zero-copy int4 path — the same #987 mechanism the f32 cache uses. 3. **Regress on footprint, not emptiness.** The #979 route probe now asserts on the accounted byte total (`== packed size`, `> 0`, `!= N*K*4`) instead of `weight_nk` staying empty (which passed straight through this bug). Adds a decline test. ## Measured (this machine) RTX 4060 laptop, 20 logical CPUs, AVX2+FMA+F16C+AVX-VNNI, **no AVX-512**. Build: `--no-default-features --features native-backend,mlas`, run `--backend native`, `DECODE_THREADS=8`, greedy. Same-binary A/B via `ONNX_GENAI_CPU_MM_MLAS_QNBIT`. Decode isolated by differencing an 88-token run against an 8-token run (min of 2). **Shared/contended host — absolute latencies are high; peak RSS and the accounting are the point and are reported together.** ### qwen05b-symzp (on-disk `total_weight_bytes` = 366,846,066 B ≈ 350 MiB) | config | admitted | peak RSS | decode ms/tok (diff) | p50 ms/tok | `resident_f32_cache_bytes` | `total_weight_bytes` | text | |---|---|---:|---:|---:|---:|---:|---| | QNBIT=0 borrowed | true | **457 MB** | 700.8 | 557.5 | **0** | 366,846,066 | ref | | QNBIT=1 MLAS | true | **1056 MB** | 544.4 | 412.6 | **246,980,608** | 613,826,674 | identical | | QNBIT=1 + `--host-ram-limit 500MB` (decline) | **false** | **438 MB** | 513.9 | 491.9 | 246,980,608 | 366,846,066 | identical | `resident_f32_cache_bytes` goes **0 → 246,980,608 (235 MiB)** when the MLAS route is taken — visible to the ledger/governor where it was 0 in *both* modes before. On decline, `admitted=false`, `total_weight` reverts to on-disk, and peak RSS + speed collapse to the borrowed path (438 vs 457 MB) — a behavioural no-op. ### qwen14b-symzp (on-disk 8,549,241,669 B ≈ 7.96 GiB) — memory measured on the branch build (accounting is deterministic, independent of the orthogonal perf commits) | config | admitted | peak RSS | `resident_f32_cache_bytes` | `total_weight_bytes` | |---|---|---:|---:|---:| | QNBIT=0 borrowed | true | ~8.0 GB (≈ on-disk; WS observed climbing to 7.95 GB) | 0 | 8,549,241,669 | | QNBIT=1 MLAS | true | **25.1 GB** | **6,995,312,640** | 15,544,554,309 | Peak **25.1 GB** reproduces the PR's ~25.5 GB, confirming the ~2x regression on this host; `resident_f32_cache_bytes` now reports **6.51 GiB** where it was 0 (`total_weight` = 7.96 GiB on-disk + 6.51 GiB packed). The 14B decode ms/tok was not isolable in reasonable wall-time on this contended shared host (a second worktree was running native inference concurrently); the memory figures are per-process and clean. **Accounted vs peak-RSS gap:** the accounted number is MLAS's own `sqnbit_packed_b_size` — the authoritative packed-B allocation the route decision already computes, and the budgetable persistent buffer (mirroring how `resident_f32_cache_bytes` accounts the primary f32 cache). Peak RSS runs higher because `SQNBitPackedB` also keeps a scales copy and packing transiently materialises a dense copy of the int4 bytes (freed afterward). ## Item 4 — can the packed buffer *replace* the mapped weights (→ ~1x instead of ~2x)? **Measured answer: no, not without a loader-level redesign.** - The weights are a **single shared read-only mmap** of `model.onnx.data` owned centrally by `WeightStore` (`onnx-runtime-loader/src/weights.rs`), kept alive via `Arc`; every initializer is a byte-range slice into that one mapping. - Kernels receive a borrowed `TensorView` per call; the CPU EP has neither per-initializer page-range info nor the authority to unmap. After MLAS packs, the MLAS route never reads that node's slice again — but other nodes on the borrowed fallback path (2-bit, int8, non-constant, or budget-declined acc0) and embeddings/lm_head/layernorm still read from the same file. - Releasing pages would need: (1) the loader to expose per-initializer page ranges, (2) a post-prepack hook after all nodes pack, (3) OS-specific discard (`DiscardVirtualMemory`/`VirtualUnlock` on Windows, `madvise(MADV_DONTNEED)` on Linux), and (4) care at page-granularity boundaries shared with borrowed-fallback initializers. - Measurement confirms coexistence today: qwen05b QNBIT=1 peak 1056 MB = borrowed 457 MB (incl. ~350 MB mmap) + ~599 MB packed side allocation; the mmap stays fully resident beside the packed buffer, so the delta is ~2x. A worthwhile follow-up, out of scope for a kernel-local change. ## Policy when the budget is too small Admitted through the exact #987 gate: `build_memory_strategy_plan` admits only when `on_disk + packed ≤ residency ceiling`. When it does not fit, `f32_weight_cache_admitted=false` → `set_mlas_sqnbit_packing_enabled(false)` → `mlas_sqnbit_owns_fp32_compute` declines → the node keeps the borrowed zero-copy int4 path. Declining is byte-identical to the borrowed path (already covered by tests), only slower per token on x86. Demonstrated above with `--host-ram-limit 500MB`. ## Tests `cargo test -p onnx-runtime-ep-cpu --features mlas --lib matmul_nbits` → **100 passed, 0 failed, 6 ignored** (includes the re-stated #979 footprint probe and the new decline test). --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624

…1056) The process-global weight-transpose cache holds one full K x N f32/f16 copy per transposed constant weight for the session. #1056 requires every session-lifetime, weight-scaled allocation to be declared to the memory plan before it is allocated, in the bytes actually allocated, and to be declinable. Reporting (TransposeCache::bytes / weight_transpose_cache_bytes) landed in 1ae696c; this finishes the governing half. - Predictor: weight_transpose_cache_predicted_bytes(&Graph) mirrors the exact kernel decision. On all platforms it counts constant-weight Gemm nodes with transB != 0 (gemm.rs transposed_b, added in #1035) at N*K*4 f32 bytes; the MatMul/FusedMatMulBias transpose call sites are cfg(macos/ios)-only, so the predictor's MatMul arm is likewise cfg-gated -- a binary predicts exactly what its own kernels allocate. The shape-keyed kernel cache instantiates a node once per activation shape, but every instance keys the global cache on (weight address, K, N), so the second instantiation hits the first entry: the transpose is held once per weight, not once per instantiation (no #1051 per-copy multiplier here). Verified by a test that runs the real Gemm kernel at prefill (m=4) and decode (m=1) and asserts predicted == bytes held, 1.00. - Admission gate: set_weight_transpose_cache_enabled(bool). When declined, MatMulPrepack::transposed_b / transposed_b_f16 return None (kernels recompute a transient transpose per call, freed each call) and cached_transpose_f32/f16 compute without inserting -- nothing session-lifetime accrues. Declining is a pure performance tradeoff: the transpose is byte-identical either way. - Wiring: engine load.rs folds the predicted bytes into resident_f32_cache_bytes and calls set_weight_transpose_cache_enabled beside the resident dequant f32 cache (#987) and MLAS SQNBit packed buffer (#1051) gates, so one plan verdict governs all three. Gates: cargo test -p onnx-runtime-ep-cpu --lib (1216 passed, 0 failed, 11 ignored); cargo clippy -p onnx-runtime-ep-cpu --lib -D warnings clean; engine --features native-backend --lib green except two pre-existing native_decode IO failures unrelated to this change (confirmed failing on origin/main). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624

#1100) ## #1056: bring `MatMulPrepack::dense` under the memory plan Fourth resident, weight-scaled buffer brought under the memory-strategy plan after the resident dequant f32 cache (#987), the MLAS SQNBit packed buffer (#1051), and the weight-transpose cache (#1079). Follows #1056's rule: *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.* ### The exact condition under which the kernel caches (verified) `MatMulPrepack::dense(index, view)` in `crates/onnx-runtime-ep-cpu/src/kernels/matmul.rs` caches a session-lifetime `Vec<f32>` iff **both**: 1. the operand is a constant initializer (`constant_inputs[index]`), and 2. `to_dense_f32_widen("MatMul", view)` returns `Cow::Owned` — i.e. the operand is **not** already a contiguous f32 view. `to_dense_f32_widen` (`crates/onnx-runtime-ep-cpu/src/dtype.rs:806`) borrows contiguous f32 zero-copy (`Cow::Borrowed`, no cache) and allocates an owned `4*numel` f32 copy for every other float case: f16/bf16/f64, or a strided/column-major f32. So a contiguous f32 constant costs **nothing** (this is why the buffer stayed invisible on the int4 and f32-contiguous models we exercise most), while an f16/bf16/f64 or non-contiguous constant costs a permanent `4*K*N` per kernel instance. ### The shape-keyed instantiation multiplier (×2) — corrected after review The executor's kernel cache is **shape-keyed** (`KernelKey { node, resolved_input_shapes }`, `crates/onnx-runtime-session/src/executor/kernel_cache.rs`). A decoder instantiates each `MatMul` node at **two** activation shapes — prefill (`m>1`) and decode (`m==1`) — as **separate** `MatMulKernel`s, each with its own `MatMulPrepack::dense`. Unlike the weight-transpose cache (process-global, keyed on the weight address, so a second instantiation reuses the first), `dense` is per-instance. The only case that populates `dense` is a constant non-f32 `B` paired with an f32 `A` (a same-half `B`+`A` takes `try_matmul_half`, and both the `m==1` GEMV and the MLAS half-prefill fast paths require `A` to be half, so they are skipped when `A` is f32). In that case **both** the prefill instance and the decode instance take the generic/direct-f32 GEMM that widens `B`, so **each retains its own `4*K*N` copy** for the session. The resident footprint is therefore `2 × 4·K·N`, not `4·K·N`. The first cut of this PR counted **one** copy — an **under**-prediction by 2×, the exact under-reporting defect #1051 corrected for the MLAS packed buffer (it reported 247 MB where steady state grew to ~592 MB, the missing factor being this same prefill/decode shape-keyed doubling). Fixed by adding `MATMUL_DENSE_DECODE_INSTANTIATIONS = 2` (mirroring #1051's `MLAS_PACKED_DECODE_INSTANTIATIONS`) and multiplying the per-node prediction by it. **Measured, not reasoned**: the ratio test below instantiates prefill + decode and counts two live copies. ### What was built - **Predictor** `matmul_dense_cache_predicted_bytes(&Graph)` — mirrors the kernel condition per node (`MatMul` + `FusedMatMulBias`, both operand indices). A graph initializer is contiguous (`WeightRef` carries only dtype + dims, no strides), so from the graph the condition reduces to *"a constant operand whose dtype is a non-f32 float"* → `4*numel`, **× `MATMUL_DENSE_DECODE_INSTANTIATIONS`** for the prefill+decode instances; a contiguous f32 constant → 0. - **`GovernedWeightCache<f32>`** — `dense: [OnceLock<Vec<f32>>; 2]` → `[GovernedWeightCache<f32>; 2]`. Declined → `dense()` widens transiently per call and frees it. The only extension needed was a read-only `filled()` accessor (reuse an existing session copy without ever running a builder); the "array of two" and the `f32` element type were already expressible, so nothing else changed in that type. - **Plan wiring** in `crates/onnx-genai-engine/src/engine/load.rs`, beside `set_resident_dequant_f32_cache_enabled` / `set_mlas_sqnbit_packing_enabled` / `set_weight_transpose_cache_enabled`, folding the prediction into `resident_f32_cache_bytes` and gating on the same `f32_weight_cache_admitted` verdict — one decision governs all four buffers. - **Decline path** — process-global `AtomicBool` written only by production (`set_matmul_dense_cache_enabled`), with a test-only thread-local RAII `DenseCacheEnabledScope` that restores on drop including on panic. No process-global mutable state tests mutate (#983/#1033/#1079). ### Predicted vs actual (ratio 1.00) — now across both instantiations Proven in one run by `predicted_dense_bytes_equal_actual_after_matmul_execution`: it computes the plan's graph prediction, then — mirroring the shape-keyed kernel cache — builds a **separate** `MatMulKernel` for the prefill shape (`m=4`) and the decode shape (`m=1`), executes each with a constant **f16** `B` and an **f32** `A`, asserts each retains a copy, and compares the **summed** `live_bytes()` against the prediction. | case | K×N | instantiations retaining a copy | per-copy | predicted (×2) | summed actual | ratio | |---|---|---|---|---|---|---| | f16 constant B, f32 A | 48×33 | 2 (prefill m=4 + decode m=1) | 6336 B | 12672 B | 12672 B | **1.00** | The test asserts `summed_actual == predicted` and `instances == MATMUL_DENSE_DECODE_INSTANTIATIONS`, so a dropped multiplier would fail it — the single-instantiation version could not. ### What happens on decline `declined_dense_cache_retains_nothing_and_is_byte_identical` runs one kernel instance admitted vs declined: | arm | bytes held after run | output | |---|---|---| | admitted | `40*24*4` = 3840 B (one copy) | reference | | declined | **0 B** | **byte-identical** | Declining is a pure performance tradeoff: the widened f32 is byte-identical whether cached or recomputed. Because the transient widen is freed at the end of each call, no session-lifetime footprint accrues when declined — and with the ×2 multiplier, admitting an f16-weight model now costs the plan the full `2×` it will actually hold. ### Peak RSS before/after — no available model populates this cache Checked every model in `C:\Users\justinchu\dev\models` for plain `MatMul` nodes with a constant operand: - `qwen2.5-14b-f32`, `qwen2.5-14b-onnx`, and all `qwen*`/`qwen05b*` variants: **0 plain MatMul nodes** — every weight routes through `MatMulNBits` (+ `GatherBlockQuantized`). - `gemma-3-27b-onnx`: exactly **1 plain MatMul**, and **both its operands are activations** (there is a `Transpose` feeding it), so `constant_inputs = [false, false]` and it caches nothing. So there is no local model that genuinely populates `dense`, and a CLI A/B would truthfully report predicted = actual = 0 — which proves nothing. Per #1056 I therefore prove predicted==actual and byte-identical-on-decline on a **synthetic graph** (the tests above) rather than reporting a meaningless zero. An f16 peak-RSS A/B awaits an export that uses plain `MatMul` with f16/non-contiguous constant weights. ### Anything not predicted exactly (over-predict, documented) The predictor **over-predicts, never under-predicts**, with documented directions: 1. **Same-half packed path**: when both operands are the same half dtype and contiguous, the node takes `try_matmul_half` and never calls `dense`, so the true cost is 0 while the predictor still counts the constant operand. Whether the *other* operand is half is not a graph-static property of the constant one, so counting is the safe direction. 2. **Prefill-only / multi-shape workloads**: a prefill-only run instantiates one copy, so the ×2 over-estimates there (safe — the gate declines sooner). Conversely a session presenting *more* than two distinct activation shapes to a node (several prompt lengths across `generate` calls) could instantiate more than two; the multiplier follows #1051's convention of bounding the autoregressive decode workload at prefill + decode, and that residual is the same documented class as #1051's. 3. **Residual gap — non-contiguous f32 constant** (a column-major weight): the kernel *would* cache it, but `WeightRef` exposes no strides, so from the graph it is indistinguishable from a contiguous f32 constant and is **not** counted — a potential under-prediction. No such operand occurs in the models this repo exercises: the one documented column-major weight (the lm_head projection, per `contiguous_b_f16`'s doc) is **f16**, so it *is* counted via the dtype rule. Documented in the predictor rather than over-counting every f32 weight (which would inflate the budget on the overwhelmingly common contiguous-f32 case). If f32 column-major MatMul weights ever appear, the loader must expose their layout for the predictor to see them. ### Gates - `cargo test -p onnx-runtime-ep-cpu --lib` — **five consecutive runs**, all `1299 passed; 0 failed; 11 ignored` (46.0 / 49.9 / 52.6 / 52.7 / 49.6 s). Includes the 4 new tests (3 in `matmul`, 1 `filled()` in `governed_weight_cache`). - `cargo test -p onnx-genai-engine --features native-backend --lib` — `496 passed; 0 failed; 1 ignored`. - `cargo clippy -p onnx-runtime-ep-cpu --lib -- -D warnings` — clean. Closes #1056. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624 --------- Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624

justinchuby force-pushed the squad/971-governed-weight-cache branch from b47787f to 32d6a07 Compare August 15, 2026 04:10

Merge remote-tracking branch 'origin/main' into squad/971-governed-we…

59d585c

…ight-cache

justinchuby merged commit 2276c12 into main Aug 15, 2026
2 checks passed

justinchuby deleted the squad/971-governed-weight-cache branch August 15, 2026 04:31

justinchuby mentioned this pull request Aug 15, 2026

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

Closed

justinchuby mentioned this pull request Aug 15, 2026

Repair #987/#979 interaction: symmetric int4 no longer caches weight_nk #1008

Merged

justinchuby mentioned this pull request Aug 17, 2026

feat(ep-cpu): bring MatMulPrepack::dense under the memory plan (#1056) #1100

Merged

This was referenced Aug 18, 2026

feat(engine): per-row device sampling routing for continuous batching (#1022) #1155

Merged

Integrate block-quantized cached-dense residency with the resource governor #666

Open

Navigation Menu

Repository navigation

Govern the CPU MatMulNBits resident f32 decode cache (#971) - #987

Govern the CPU MatMulNBits resident f32 decode cache (#971)#987
justinchuby merged 2 commits into
mainfrom
squad/971-governed-weight-cache

justinchuby commented Aug 15, 2026

Uh oh!

codecov Bot commented Aug 15, 2026 •

edited

Loading

Uh oh!

justinchuby commented Aug 15, 2026

Uh oh!

Uh oh!

github-actions Bot commented Aug 15, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

justinchuby commented Aug 15, 2026

Summary

Part 1 — truthful accounting

Part 2 — governed decision

Acceptance evidence

Gates (verbatim)

Uh oh!

codecov Bot commented Aug 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

justinchuby commented Aug 15, 2026

Re: correction on the real discriminator (issuecomment-5299053308)

Uh oh!

Uh oh!

github-actions Bot commented Aug 15, 2026

🔴 Benchmark Regression Detected

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Aug 15, 2026 •

edited

Loading