Evict MLAS SQNBit packs by owner, and key them on every operand (#1735)

[justinchuby]

Closes #1735.

The defect

MlasPackedCaches keys the shared SQNBit packs on the operands' addresses and holds no claim on the memory there. An address names a weight only while that weight is live: once freed, the next allocation of matching shape and pack parameters is served the previous weight's packed bytes. The route counter still increments, because the right route did run — it just multiplied another weight's rows. This is #1726's defect in the packed store rather than the transpose caches, and clear_mlas_packed_caches() at executor teardown only closed the cross-model case.

Two gaps, closed together

Lifetime. get_or_build_shards/get_or_build_packed now report whether this call installed the entry, and only the installing kernel records ownership. A retiring kernel evicts what it installed.

Ownership is install-only on purpose: a kernel that merely hit an entry must leave it alone, or a retiring prefill instance would pull the pack out from under a live decode sibling and force a re-pack — the cost the shared store (#1056) exists to remove. This is the defect Opus caught in #1732's first draft, so it is asserted directly.

Identity (not in the issue). A CompInt8 pack bakes the scales and zero points into its per-block sums, but the key covered only the quantized weight — two nodes sharing a B initializer under different scales would serve each other packs. The key now carries all three operand addresses, and a weight whose scales have no stable host address is treated as unshareable rather than shared under a partial identity. This is the "stable full weight identity where available" the directive asked for, and it also strengthens the recycled-address case, since a colliding address must now match on scales too.

Bonus. clear_mlas_packed_caches() drained MLAS_PACKED_GLOBAL directly while cfg(test) traffic goes to the thread-local store — a silent no-op in every test that called it. Now routed through with_mlas_packed_caches.

Design notes

MlasPackedOwnership exists so the Drop does not land on MatMulNBitsKernel itself: a type that implements Drop cannot be moved out of, which makes the ..test_kernel(..) functional-update syntax used throughout the suite illegal (70 compile errors on the first attempt). Moving out of a struct whose fields have Drop is fine.

Production teardown ordering is unchanged and now guarded: clear_weight_transpose_caches() and clear_mlas_packed_caches() run in Executor::drop before self.buffers.drain().

Tests

Every test was verified to fail under its own mutation and no other:

test	mutation that kills it
a_recycled_weight_address_must_not_serve_the_previous_weights_mlas_pack	Drop made a no-op → round 1 serves weight 0's output for weight 1
a_shared_mlas_pack_outlives_the_sibling_that_only_shared_it	ownership recorded on hit too → sharer claims the entry
one_quantized_blob_under_two_scale_tensors_must_not_share_a_pack	scales_addr dropped from the key → second scales tensor gets the first's pack
the_packed_store_stays_consistent_under_concurrent_install_and_evict	(drives MlasPackedCaches directly; the per-test store is thread-local, so kernels on different threads cannot contend)
weight_derived_caches_are_cleared_before_their_buffers_are_freed	clears moved after buffers.drain()

The falsifier runs 8 alternating rounds and requires a cross-weight address reuse before it will pass — "some address was reused" is satisfiable by a weight landing on its own former address, which proves nothing. Eight rounds are needed because glibc first serves the ~2 MiB pack by mmap; freeing one raises the dynamic mmap threshold so later same-size requests come from the brk heap, where addresses are reused. Reference outputs are taken against a drained store, so a stale hit cannot poison the reference and make the comparison a tautology.

Routing note: both kernel-level tests drive try_mlas_sqnbit under NXRT_CPU_GEMM_BACKEND=mlas at m=4. On x86_64 the native acc4 int4 route wins before MLAS is consulted, and MLAS refuses asymmetric CompInt8 at M=1 on hosts without a correct AVX2 kernel — either would have made the tests silently vacuous.

Validation

./roy_validate.sh — PASS=21 FAIL=0 SKIP=0, including H check-win-arm64 (cargo-xwin) and G2 test-aarch64-qemu (the aarch64 suite actually executed, not just type-checked).

• default (no-mlas) build/test/clippy -D warnings, --no-default-features, --all-features
• K no-mlas-artifacts — the default zero-MLAS artifact is untouched; the research-only feature policy is preserved
• concurrency: RUST_TEST_THREADS=32 × 6 clean runs, 1621 lib tests each
• Miri (-Zmiri-tree-borrows): passes on the_packed_store_stays_consistent_under_concurrent_install_and_evict. Miri cannot execute the other three — they call into MLAS through FFI. Coverage here is the safe store/ownership logic only, and that limit is stated rather than papered over.

The default (no-mlas) lane caught a real cfg regression during development: an inserted field consumed the #[cfg(feature = "mlas")] that belonged to mlas_shards at both initializer sites.

⚠️ The New weight-derived caches must be governed gate is expected to fire — it matches on diff text, so a modified field declaration reads as a new cache. No new cache is introduced; MlasPackedOwnership holds keys, not buffers.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.06%. Comparing base (c498792) to head (951bd31).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1738      +/-   ##
==========================================
+ Coverage   79.75%   80.06%   +0.30%     
==========================================
  Files         408      408              
  Lines      193197   192932     -265     
  Branches   193197   192932     -265     
==========================================
+ Hits       154088   154468     +380     
+ Misses      33759    33109     -650     
- Partials     5350     5355       +5     

Flag	Coverage Δ
cli-ort-linux	72.47% <ø> (ø)
cli-ort-windows	72.06% <ø> (ø)
mlas	85.33% <ø> (+0.13%)	⬆️
offline	80.18% <100.00%> (+0.31%)	⬆️

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	82.26% <100.00%> (+2.52%)	⬆️

... and 18 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[justinchuby]

Independent Opus review — disposition

Reviewed adversarially against six attack surfaces. Five came back unfounded after the reviewer traced the concrete interleavings; both actionable findings were test-robustness issues and are fixed in bfb6b1e-range follow-up commit.

Unfounded, with the reasoning worth recording:

• Install-only ownership is complete. The worry was that one kernel could OnceLock::set mlas_owned.packed twice under different keys, silently dropping the second and leaking that entry forever. It cannot: shared_mlas_packed and shared_mlas_prepacked are mutually exclusive (self.weight_prepacked returns before the non-prepacked path), each is memoized behind self.mlas_packed, and comp is a pure function of self.accuracy_level so it cannot vary between calls. A kernel can install both a shards entry and a packed entry — those are different OnceLocks into different maps, and each remove_* hits its own.
• remove_* vs a concurrent get_or_build_* is perf-only. Two kernels can never both own one entry: installed=true requires a Vacant entry under the mutex, and the prior owner removes before a later kernel can re-install. If an owner evicts an entry a sharer is using, the sharer already memoized its Arc and never consults the store again; a third kernel merely re-packs, and always from its own operands at its own key's addresses, never from foreign memory.
• The extended key does not break prefill/decode sharing. Every store path requires can_prepack, which requires the scales and zero points to be constant initializers — delivered as views into the same shared buffers for both instances. The key is taken from the raw inputs before any densification, so that cannot desync the addresses either.
• Drop order inside the kernel is not observable. mlas_shards/mlas_packed drop before mlas_owned, but both the memo and the store hold Arcs of the same data; refcounting means eviction never frees memory another holder references, in either order.
• contiguous_host_addr not using contiguous_host_slice::<u8> is correct (that helper forms a numel-length u8 slice, understating a wider dtype), and None → treated as unshareable is what every path actually does — there is no partial-key fallback.

Fixed:

• The vacuousness guard tracked only the packed weight's address while a stale hit now needs all three to collide, so on an allocator whose size classes don't recycle in lockstep the test could report recycled=true without the key ever matching — passing even with eviction reverted. Now tracks the whole key tuple.
• More importantly the falsifier no longer leans on the allocator: each round asserts the store is empty before it starts, which the retiring previous kernel guarantees on every allocator. Re-verified — a no-op Drop is now killed deterministically at round 1 (round 1: the retired kernel left its pack in the store) instead of waiting for an address to be reused.
• The final recycling check is no longer an assertion, because address reuse is the allocator's choice, not the code's: under ASan/valgrind/some musl and macOS configurations a correct implementation would have failed it. Reported loudly instead, so a skipped arm can't be misread as an exercised one. On this glibc host both arms run.

Noted, not changed: the env-var save/restore in the two new tests is not panic-safe — an assertion failing between set_var and the restore leaves NXRT_CPU_GEMM_BACKEND set and poisons backend_env_lock for every later test. That is the established pattern across ~8 pre-existing tests in this file, so the new tests conform rather than diverge. Fixing it properly means a scope guard for all of them; out of scope here, worth its own issue.

Final validation

./roy_validate.sh — PASS=21 FAIL=0 SKIP=0, re-run on the final tree.

Includes H check-win-arm64 (cargo-xwin, aarch64-pc-windows-msvc) and G2 test-aarch64-qemu — the aarch64 suite executed under qemu-user, not merely type-checked. No step was skipped for a missing toolchain, so the 21 is a complete matrix rather than a narrowed one.

Plus: RUST_TEST_THREADS=32 × 6 clean runs (1621 lib tests each), and Miri -Zmiri-tree-borrows on the store's concurrency test. Miri cannot execute the other three — they cross into MLAS via FFI — so Miri coverage here is the safe store/ownership logic only.

[justinchuby]

New weight-derived caches must be governed — judgement

The gate matched these two added declarations:

struct MlasPackedOwnership {
    shards: OnceLock<MlasPackedKey>,
    packed: OnceLock<MlasPackedKey>,
}

They match the pattern OnceLock *<.*Packed, but they are not caches and hold no buffer. MlasPackedKey is a plain Copy struct of usize/bool/enum fields — three operand addresses, n, k, bits, block_size, has_zero_points, and the compute type. Two of them is on the order of 80 bytes, fixed, independent of weight size, and it is per kernel instance rather than per process.

Its purpose is the opposite of a cache: it is the record of what this kernel installed so that its Drop can remove those entries from the shared store. It shrinks long-lived memory rather than adding any — before this PR an entry lived until executor teardown; now it is evicted when its installing kernel retires.

No new weight-derived buffer is introduced anywhere in this PR. The buffers involved (MlasPackedCaches.shards / .packed) already existed and are already accounted — SQNBIT_PACKED_LIVE_BYTES in mlas-sys counts each packed buffer once, and int4_acc0_mlas_packed_accounting_equals_actual_allocated still passes. Their byte total can only go down as a result of this change.

Recording that judgement with the weight-cache-reviewed label, per the gate's own instructions.

For completeness, the accounting is untouched because the store's contents are unchanged — the same Arc<Vec<Option<MlasShard>>> and Arc<MlasPreparedPacked> values, under a key with two more usize fields. The only lifetime change is that entries now leave the store earlier.

[justinchuby]

Merge record

Merging under Justin's standing direct-merge directive (2026-08-19T14:20:27: comprehensive local validation including cross-platform, then merge directly rather than waiting on queued Actions).

Local matrix: PASS=21 FAIL=0 SKIP=0, re-run on the final tree after the review hardening. No step was skipped for a missing toolchain, so the 21 is the complete matrix — including H check-win-arm64 (cargo-xwin → aarch64-pc-windows-msvc) and G2 test-aarch64-qemu, which executes the aarch64 suite under qemu-user rather than merely type-checking it.

Required-CI position, stated honestly. One check is red: Rust coverage (macOS arm64). I checked it is not mine before merging — it fails identically on unmodified main (runs 32565237825 and 32564899444, alongside the rest of the chronic set from #1600: Rust (Windows ARM64), CLI ORT (Windows x86_64), Rust coverage (Windows x86_64)). It is not attributable to this PR, and this change's aarch64 behaviour is covered locally by G2 executing the suite on aarch64. Every other completed check passes, including New weight-derived caches must be governed after the judgement label.

The remaining checks are queued, not failing. Per the directive I am not waiting on them; if any produces a real code failure afterwards I will fix it on main rather than leave it.