Repository navigation
Share one MLAS SQNBit packed buffer per weight (#1056) - #1098
Conversation
The executor's kernel cache is shape-keyed, so an autoregressive decoder compiles two MatMulNBits kernel instances per node -- one for prefill (m > 1) and one for decode (m == 1). Before this change each instance packed its own full copy of the same constant weight into its own mlas_shards OnceLock, so the resident MLAS SQNBit packed footprint was 2x the single-copy cost (measured on qwen05b-symzp: 169 packs on a 1-token run, 338 on a multi-token run). Introduce a process-global, weight-identity-keyed store (MlasPackedCaches) keyed on (address, N, K, bits, block_size, has_zero_points, compute_type). The first kernel instance to reach a weight packs it once; the sibling instance takes the same Arc. The session now holds one packed copy per weight. clear_mlas_packed_caches() runs on Executor drop -- the same lifetime boundary as weight_transpose::clear_all -- to close the recycled-address hazard (#845/#1079). Accounting is updated in lockstep: MLAS_PACKED_DECODE_INSTANTIATIONS goes from 2 to 1, so resident_dequant_f32_cache_bytes (the memory plan's prediction) equals mlas_sqnbit_packed_live_bytes (the actual allocation), ratio 1.00. The accounting test now also asserts pointer identity of the shared Arc across the prefill and decode instances. Test isolation uses a cfg(test) thread-local store so each libtest thread gets a private cache (no cross-test recycled-address contamination) while a single test's prefill+decode still share; production reads only the global store (no new global mutable state tests mutate, #983/#1033/#1079). Refs #1027, #1051, #1056. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
|
Closing per owner direction: we are not bundling MLAS by default and will instead absorb its optimizations into our native kernels. Deduplicating MLAS's packed buffer would be optimizing the wrong thing. Superseded by an investigation into where MLAS SQNBit's 1.25x decode advantage comes from and how much is reachable without a resident copy. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1098 +/- ##
==========================================
- Coverage 80.48% 79.91% -0.57%
==========================================
Files 367 369 +2
Lines 157217 160014 +2797
Branches 157217 160014 +2797
==========================================
+ Hits 126531 127879 +1348
- Misses 25974 27410 +1436
- Partials 4712 4725 +13
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
🔴 Benchmark Regression DetectedComparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).
Visual flags: Host infoWhat this cannot catch
|
|
Merging, with a note on how this now fits. Validated after merging current Context change while this was in flight. The owner set the direction that we do not bundle MLAS by default and instead absorb its optimizations natively (#1091), and #1104 has now done exactly that for this route: the int4 decode advantage turned out to be register/N-blocking rather than layout, so it ported into our borrowed path with no packed buffer at all -- 1.56x on the 14B at unchanged peak RSS. That reframes this PR rather than cancelling it. MLAS stays in the tree as a measurement reference, because keeping both paths in one binary is what made every one of these gaps measurable. A reference that costs half as much and accounts itself correctly is worth having, and leaving a validated branch to rot is worse than landing it. Merging as maintenance of the reference, not as investment in the dependency. The honest correction in your report is the most valuable part. My hypothesis when I dispatched this was that halving the packed cost would flip the 14B from declined to admitted. You measured that it does not: admission compares the expanded footprint (weights + cache) against the ceiling, and the dedup took that from 26.47 GB to 17.51 GB against a 17.13 GB ceiling -- missing by 378 MB, 2.2%. Reporting "the thing you predicted would happen does not happen, and here is the arithmetic" is exactly right, and the derived figure -- the admit threshold dropping from 0.386 to 0.256 of RAM -- is more useful than the yes/no I asked for. The rest reproduces as described: pack count 338 → 169 on a multi-token run, now equal to the 169 of a single-token run; accounted 632,887,808 → 316,443,904 with ratio 1.00 at both scales; the accounting constant moved 2 → 1 in the same commit as the allocation change, which is what keeps predictor and allocation from drifting. One detail worth carrying to the remaining per-instance caches: keying the store on weight identity and clearing it at |
Share one MLAS SQNBit packed buffer per weight (#1056)
Refs #1027 (MLAS SQNBit route), #1051 (packed-buffer accounting), #1056 (this dedup).
The problem
MatMulNBitsint4accuracy_level=0nodes route to MLAS SQNBit CompFp32 (#1027). Theexecutor's
KernelCacheis shape-keyed, so each node compiles two kernel instances —prefill (
m > 1) and decode (m == 1) — and before this change each instance packed itsown full copy of the same constant weight. The resident packed footprint was therefore
2xthe single-copy cost, held for the whole session.
The fix
A process-global, weight-identity-keyed store (
MlasPackedCaches) keyed on(address, N, K, bits, block_size, has_zero_points, compute_type). The first kernel instance toreach a weight packs it once; the sibling instance takes the same
Arc. The session now holdsone packed copy per weight.
different-shape weight (allocator recycling a freed address — the onnx-runtime-ep-cpu: weight-transpose cache keyed only by pointer serves wrong-length transposes on address reuse #845/Govern the weight-transpose cache under the memory plan (#1056 item 2) #1079 hazard) misses
rather than serving the wrong bytes.
clear_mlas_packed_caches()runs onExecutordrop — the same lifetime boundary asweight_transpose::clear_all— closing the same-address/same-shape/across-lifetimes window.MLAS_PACKED_DECODE_INSTANTIATIONSgoes2 -> 1,so
resident_dequant_f32_cache_bytes(the plan's prediction) equalsmlas_sqnbit_packed_live_bytes(the actual allocation). The existing accounting test additionally asserts pointer identity
of the shared
Arcacross the prefill and decode instances.only the global store; tests use a
cfg(test)thread-local store (each libtest thread gets aprivate cache, so no cross-test recycled-address contamination, while a single test's prefill+decode
still share).
Acceptance criteria
1. Predicted bytes == actual bytes, ratio 1.00.
resident_f32_cache_bytes(plan prediction) vslive_total(profiler actualSQNBIT_PACKED_LIVE_BYTES), measured on-model:The dedup test (
int4_acc0_mlas_packed_accounting_equals_actual_allocated) that ties the plan'sprediction to the profiler's actual bytes stays green and now also asserts the shared
Arc.2. Pack count halves.
ONNX_GENAI_PROFILE_MM=1,[mm_prepack] calls=onqwen05b-symzp(169 weight boundaries):
Before, the multi-token run packed twice as many buffers as the 1-token run; after, they pack the
same count. (The 1-token run already packed 169 before, but the pre-dedup predictor still
accounted
2x = 632,887,808for it — over-report, the safe direction; after, both the pack countand the accounting are single-copy.)
3. Peak RSS + accounted, with ratios, both models, before/after. Every number measured on this
host (Windows, 68,535,443,456 B RAM, CPU-only, AVX2/FMA/F16C/AVX-VNNI). Peak RSS = polled
PeakWorkingSet64while running; CPU time =TotalProcessorTime.qwen05b-symzp (weights 366,846,066 B), multi-token autoregressive run:
QNBIT=0)The packed accounting halved (632,887,808 -> 316,443,904) and peak RSS dropped 317.4 MB — almost
exactly the one deduplicated packed buffer (316,443,904 B).
qwen14b-symzp (weights 8,549,241,669 B). Default residency ceiling =
0.25 x RAM=17,133,860,864 B. Admission tests the expanded footprint (
on-disk weights + packed cache):Does the 14B flip declined -> admitted at the default ceiling? No — but only just, and the reason
is precise. The dedup halved the predicted packed cache (17,925,488,640 -> 8,962,744,320) and
shrank the expanded footprint from 26,474,730,309 to 17,511,985,989. But admission compares that
expanded footprint (weights + cache), not the cache alone, against the
0.25 x RAMceilingof 17,133,860,864 B. After the dedup the expanded footprint is 17,511,985,989 B — still 378,125,125 B
(2.2%) over the default ceiling, so it stays declined and runs the borrowed zero-copy path
(peak ~8.6 GB, unchanged from before).
What the dedup does change is the admission threshold: admitting the 14B previously required a
ceiling >= 26,474,730,309 B = 0.386 of RAM; it now requires >= 17,511,985,989 B = 0.256 of RAM
— i.e. barely above the 0.25 default.
--host-ram-limit 18GiB(0.263) now admits it at peak17,606.9 MB, comfortably inside 68.5 GB. So the dedup moves the 14B from "unreachable without
allowing a 26.5 GB expansion" to "one notch above the default," but does not cross the 0.25 line
on its own on this box. (The original prediction that it would flip rested on comparing the ~8.4 GB
cache to the ceiling; the gate actually tests the 17.5 GB expanded footprint.)
4. Byte-identical generated text (SHA-256 of generated text), greedy decode:
EF7CA14F…EF7CA14F…FD4972FF…FD4972FF…EB88829D…EB88829D…EB88829D…Identical across before/after and across admitted/declined (the borrowed zero-copy path and the MLAS
packed path produce the same tokens). qwen05b route ON and route OFF also match (
EF7CA14F…).5. No new process-global mutable state that tests mutate. Production writes only the
LazyLockglobal; tests use acfg(test)thread-local store restored automatically by eachlibtest thread ending. No env/RAII toggles were added to production.
Gates
cargo test -p onnx-runtime-ep-cpu --features mlas --lib matmul_nbits— 110 passed, 0 failed, 6 ignored.cargo test -p onnx-runtime-ep-cpu --lib— five consecutive runs:1269/0/11,1269/0/11,1269/0/11,1269/0/11,1269/0/11(passed / failed / ignored).cargo clippy -p onnx-runtime-ep-cpu --lib -- -D warnings— clean (both default and--features mlas).Cases not deduplicated
None on the constant-weight route. Every MLAS SQNBit route (
weight_prepacked, static shards, andthe
NO_SHARDA/B) goes through the shared store. The only non-shared fallback is a weight with nostable contiguous host address to key on — which never occurs on the constant-weight (
can_prepack)route this touches, since the initializer is a contiguous mmap slice. A non-constant weight rebuilds
a transient pack per call and retains nothing, so there is nothing to share.
Note:
Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624.