Repository navigation
MatMulNBits: stop the borrowed int4 path pre-empting MLAS SQNBit CompFp32 (up to 33x, 78x->2x vs ORT) - #1027
Merged
Merged
Conversation
…Fp32 `Kernel::execute` evaluates the `bits == 4 && accuracy_level == 0` borrowed zero-copy branch *before* `try_mlas_sqnbit`. That ordering made the accuracy_level=0 -> MLAS SQNBit CompFp32 route -- which `try_mlas_sqnbit`'s own doc comment describes as the intended path, "matching how ORT/onnxruntime-genai run those models" -- unreachable dead code. The consequence on x86_64 is severe. `borrowed_affine_int4_matmul`'s vectorized block dot is `#[cfg(target_arch = "aarch64")]`, so on x86 every `bits=4, accuracy_level=0` node ran a scalar nibble-unpack GEMV, and for m > 1 it re-streamed the entire weight once per activation row with no GEMM blocking. accuracy_level=0 is exactly what Foundry `cuda-gpu` int4 exports emit, so this is the common case, not a corner. Measured on an AMD EPYC 9V74 (AVX2/FMA/F16C, no AVX-512), ORT 1.27 CPU EP, both runtimes pinned to 8 intra-op threads, interleaved A/B, p50 of 9 runs after 3 warmups: K=1024 N=3072 M=1 block32 1.587 ms -> 0.169 ms (9.4x) vs ORT 0.097 ms K=3584 N=3584 M=1 block32 4.625 ms -> 0.525 ms (8.8x) vs ORT 0.341 ms K=1024 N=3072 M=16 block32 18.118 ms -> 0.959 ms (18.9x) vs ORT 0.344 ms K=1024 N=3072 M=128 block32 116.348 ms -> 4.619 ms (25.2x) vs ORT 1.746 ms K=3584 N=3584 M=128 block32 385.308 ms -> 13.139 ms (29.3x) vs ORT 7.121 ms versus ORT the same node moves from 12x-78x slower to 1.5x-4.5x slower. Parity against ORT becomes bit-identical (max_abs = 0.0) because both runtimes now execute the same MLAS kernels. The setup cost is not hidden: the one-time MLAS shard pack for the K=3584 N=3584 block-32 weight is 55.8 ms for 6.4 MB (`ONNX_GENAI_PROFILE_MM=1`, `phase=mlas-shards calls=1 this=55.80ms`), against a per-call saving of ~4.1 ms, so it repays itself in ~14 calls and is amortized over the session. `mlas_sqnbit_owns_fp32_compute` states the ownership rule explicitly rather than reordering branches: * MLAS takes the node only when `can_prepack` -- B/scales/zero-points are graph constants -- so the packed buffer is built once per session. With dynamic weights MLAS would repack on every call, and the borrowed zero-copy path keeps those. * `sqnbit_packed_b_size` is MLAS's own "do I have a kernel for this shape on this host" probe; when it says no, the borrowed path stays the fallback. * Without the `mlas` feature the predicate is `false`, so non-MLAS builds are byte-for-byte unchanged. #979's memory intent is preserved: MLAS's packed buffer is int4-sized, so the resident f32 `weight_nk` expansion (~8x the file size) is still never built. The #979 regression tests now assert that invariant directly (`weight_nk` stays empty) plus the route actually taken, derived from the same predicate production code uses, so they stay exact on hosts where MLAS declines instead of degrading into tautologies. Route proof under `mlas` uses per-kernel state rather than the process-global counters, which are shared with tests running in parallel. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Review follow-up. `Int4Acc0RouteProbe::assert_fast_route` asks `mlas_sqnbit_owns_fp32_compute` what should have happened, which catches "the code disagrees with its own policy" but cannot catch "the policy itself regressed to never choosing MLAS" -- change the predicate to return `false` and those assertions quietly follow it. `int4_acc0_constant_weight_reaches_mlas_on_x86_64` pins the concrete case this fix exists for, with a hardcoded expectation: on x86_64 with the vendored MLAS, a constant-weight symmetric int4 `accuracy_level = 0` node at block size 32 must end up on MLAS SQNBit. Reintroducing the branch order that made that route dead code now fails a test even if the predicate is edited to agree with it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1027 +/- ##
==========================================
+ Coverage 78.62% 78.96% +0.34%
==========================================
Files 352 366 +14
Lines 142679 150463 +7784
Branches 142679 150463 +7784
==========================================
+ Hits 112178 118818 +6640
- Misses 26071 27001 +930
- Partials 4430 4644 +214
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
|
Owner
Author
|
Reproduced the speed-up independently, and it is real. But the memory claim does not hold, and on a large model the regression is severe enough that this must not land default-on as written. Requesting changes. Host: RTX 4060 laptop, 20 logical CPUs, 68.5 GB RAM. The speed-up reproduces
2.68x on decode, and the generated text is byte-identical between the two. The one-time cost is visible and matches your report: +1.05 s fixed, repaid after ~11 tokens. The memory claim does not hold"#979 is preserved" is argued from
It does not scale with thread count (986 / 1027 / 981 / 1061 MB at 2 / 4 / 8 / 16 threads), so these are not per-thread shards -- it is a second, private copy of the weights held beside the still-resident mapped file, packed to roughly 2x the int4 bytes. Meanwhile For context on the cost: #979/#989 took the 14B from ~66 GB to 7.80 GB, which is what makes it runnable on a 32 GB machine. 25.5 GB gives most of that back. On the 14B the MLAS route was also slower in wall time for a short run -- 115.4 s against 72.0 s for 8 tokens -- because the pack is paid before any token is produced. What I am asking forThe diagnosis is right and the win is worth having. The gate is just missing a dimension: it asks "does MLAS have a kernel for this shape" and never asks "can we afford the packed buffer".
Not merging until 1-3 are in. Happy to be shown the packed buffer can replace the mapped weights rather than sit beside them -- if the borrowed path can be dropped once packing is done, the delta becomes ~1x instead of ~2x and the calculus changes. |
justinchuby
pushed a commit
that referenced
this pull request
Aug 15, 2026
…ht copy Measured while validating #1027. Routing int4 MatMulNBits to MLAS SQNBit is a 2.68x decode win on the 0.5B (146 -> 54.4 ms/token, byte-identical output) and simultaneously takes peak working set from 8.17 GB to 25.5 GB on the 14B, because the packed buffer is a private ~2x-int4-sized copy held beside the still-resident mapped weights rather than replacing them. Flat across 2/4/8/16 threads, so it is not per-thread sharding. Recorded next to the f32 weight cache because it is the same defect twice: `resident_f32_cache_bytes` reported 0 in both modes, and the PR argued the #979 invariant held by asserting `weight_nk` stays empty -- a proxy that a different allocation of the same bytes passes straight through. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
marked this pull request as ready for review
August 15, 2026 23:52
justinchuby
added a commit
that referenced
this pull request
Aug 16, 2026
…n int8 dot product (up to 15x) (#1028) ## Summary `try_mlas_sqnbit` short-circuits `m < sqnbit_decode_min()` back to the hand int4/int8 decode kernels, on the premise that they beat MLAS SQNBit CompInt8 at small `m` while also avoiding MLAS's one-time packing. That premise is a claim about the **host**, but it was encoded as a constant: ```rust let hand_decode_is_fast = self.bits == 4 && self.accuracy_level == 4 && !prefer_arm64_mlas_decode; ``` The hand kernels are only fast where the int8 accumulation is a single instruction. On x86_64 that means AVX-VNNI or AVX-512-VNNI (`vpdpbusd`) — `DotKernel::supports_int4_direct` already requires VNNI for the int4-direct route, and without it the AVX2 fallback emulates the dot with `vpmaddubsw` + `vpmaddwd` + widening adds. On such a host the short-circuit sent **every M=1 `accuracy_level = 4` decode** to a kernel an order of magnitude slower than the MLAS SQNBit kernel sitting directly behind the gate. This is the decode hot path: M=1 is one generated token. ## Measurement Single-node `com.microsoft::MatMulNBits` models run through **both** our native CPU EP and the **real ORT 1.27 CPU EP** in one process, interleaved, via `bench_generic`. - Host: AMD EPYC 9V74 (Azure), 32 vCPU / 16 physical cores, **AVX2 + FMA + F16C — no AVX-512, no AVX-VNNI, no AMX**. - **Thread-matched**: `--ort-intra-threads 8` vs `ONNX_GENAI_CPU_DECODE_THREADS=8`. - p50 of 15 runs after 5 warmups; before/after differ only by this commit. -⚠️ Shared, contended runner (load average ~17). Absolute times are inflated; ratios are interleaved so they hold. Treat <1.2x as noise. | model (int4, symmetric) | before | after | ORT | before/ORT | after/ORT | gain | |---|---|---|---|---|---|---| | K=896 N=4864 M=1 block32 | 1.416 ms | **0.094 ms** | 0.051 ms | 21.6x | 1.85x | **15.1x** | | K=1024 N=3072 M=1 block16 | 1.074 ms | **0.088 ms** | 0.045 ms | 18.7x | 1.96x | **12.2x** | | K=1024 N=3072 M=1 block32 | 0.762 ms | **0.078 ms** | 0.041 ms | 16.1x | 1.89x | **9.8x** | | K=1024 N=3072 M=1 block128 | 0.570 ms | **0.073 ms** | 0.038 ms | 11.8x | 1.91x | **7.8x** | | K=3072 N=1024 M=1 block32 | 0.734 ms | **0.101 ms** | 0.042 ms | 14.6x | 2.42x | **7.3x** | | K=3584 N=3584 M=1 block32 | 1.386 ms | **0.191 ms** | 0.112 ms | 10.7x | 1.71x | **7.3x** | | K=1024 N=3072 M=1 block64 | 0.499 ms | **0.074 ms** | 0.038 ms | 10.8x | 1.94x | **6.7x** | **Against ORT these nodes move from 10.7x–21.6x slower to 1.7x–2.4x slower**, parity `PASS` on every one. `m >= sqnbit_decode_min()` (M=16, M=128) is unaffected — it was already above the crossover and measures flat (1.0x–1.4x, within noise on this host). ## Design `hand_int8_decode_has_native_dot()` reads the **selected** `DotKernel` rather than re-probing CPUID, so `ONNX_GENAI_CPU_DOT_KERNEL`-style overrides and the test harness stay consistent with what actually executes. - **x86_64**: requires `uses_vnni_int4_direct()` (AVX-VNNI / AVX-512-VNNI). - **aarch64**: unconditionally `true`. NEON is baseline on ARM64, so `Neon`/`NeonDot` always have a real dot product — **ARM behaviour is unchanged**. - **other**: `false`. The scalar `DotKernel` has no dot product, so MLAS (when present) should take the node. On a VNNI x86 host the gate evaluates exactly as before, so this is a strict improvement for AVX2-only hardware and a no-op everywhere else. ## Deliberately *not* changed — measured losses that remain | region | after | why declined | |---|---|---| | int4 **asymmetric** M=1 | ~12.4x slower than ORT | MLAS's `SQ4BitGemmM1Kernel_CompInt8_avx2` asymmetric kernel is numerically broken (~46% disagreement, already refused in-tree by `host_supports_mlas_sqnbit_m1_asym_int8`). Correctness wins. Serving these with CompFp32 instead would need a **second** packed weight cache, since `SQNBitPackedB` carries its compute type and the session cache holds exactly one — that memory cost needs its own justification, so it is not bundled here. | | `bits = 8` M=1 | 13.8x–16.2x slower | MLAS SQNBit has **no 8-bit x86 kernel**; there is nothing faster to defer to. Needs its own kernel work. | Both are reported rather than papered over: this PR does not claim to fix them. ## Test changes Two existing tests encoded the old unconditional premise: - `matmulnbits_try_mlas_gates_decode_by_m_threshold` — now expects `Some(())` below the crossover where the host has no native dot, and `None` where it does. The decode/prefill split is still regression-locked; the *decode* half is now host-conditional, which is the actual invariant. - `matmulnbits_accuracy4_prepack_reuses_selected_weight_format` — the invariant under test ("one weight format, chosen once, reused across calls, never the f32 expansion") is unchanged; only *which* cache legitimately owns it differs by host, so the assertion now branches on the same predicate. New: `hand_int8_decode_native_dot_matches_selected_kernel` pins the predicate to `selected_dot_kernel()` per architecture. ## Validation ``` $ cargo test -p onnx-runtime-ep-cpu --features mlas --release test result: ok. 1076 passed; 0 failed; 15 ignored $ cargo test -p onnx-runtime-ep-cpu --release # no mlas test result: ok. 1079 passed; 0 failed; 10 ignored $ cargo fmt --all # clean $ cargo clippy -p onnx-runtime-ep-cpu --features mlas --all-targets --release -- -D warnings 0 errors $ python3 scripts/check_feature_gate_coverage.py ✓ All feature-gated fast paths have fallback coverage. ``` ## Relationship to #1027 Independent. #1027 fixes `accuracy_level = 0` (borrowed path pre-empting MLAS CompFp32); this fixes `accuracy_level = 4` M=1 (decode crossover assuming an ISA the host may not have). They touch different gates and can land in either order. ## Review Independent Rubber Duck review (Opus, read-only): **APPROVE**, with two items I acted on. **1. Dynamic weights (raised by the reviewer, dismissed there, kept and fixed here).** The reviewer flagged and then dismissed the possibility that routing decode to MLAS regresses non-constant weights. I did not dismiss it, because it is measurable: with `can_prepack == false` there is no session-lifetime buffer to hold MLAS's packed weight, so MLAS repacks the **entire weight on every call** — measured at **55.8 ms for a 6.4 MB int4 weight** (`ONNX_GENAI_PROFILE_MM=1`) against a sub-millisecond hand decode. That is a large regression on exactly the case the old constant `false` gate happened to protect. **Fixed** in `bef29131b`: the gate is now `!can_prepack || hand_int8_decode_has_native_dot()`. Dynamic weights keep the hand path on every ISA — a slow kernel beats repacking megabytes per token. `matmulnbits_accuracy4_dynamic_weight_decode_keeps_hand_path` locks that in, and `matmulnbits_try_mlas_gates_decode_by_m_threshold` now declares its weights constant, which is the only case where the ISA question is live. **2. Tautological test (MINOR).** `hand_int8_decode_native_dot_matches_selected_kernel` cross-checked the predicate against `selected_dot_kernel().uses_vnni_int4_direct()` — the implementation restated. **Fixed** in the same commit: it now checks CPUID directly (`avx512vnni || avxvnni`), and skips when `ONNX_GENAI_CPU_DOT_KERNEL` overrides the selection, since there the predicate must follow what actually executes rather than what the hardware advertises. ## CI `main` is currently red for reasons that predate and are untouched by this PR — verified by checking out unmodified `origin/main` in this worktree, where `cargo fmt --all -- --check` already flags `crates/onnx-genai-ort/src/lib.rs` and `crates/onnx-runtime-ep-cuda/src/kernels/matmul_nbits.rs`. The last 8 `CI` runs on `main` all conclude `failure` (independent `CLI ORT` build, `CUDA compile` inventory/clippy and `Rust (Windows ARM64)` test failures). This PR changes exactly one file, `crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs`, and the fmt-flagged set is byte-identical before and after. Locally green on this branch: ``` cargo test -p onnx-runtime-ep-cpu --features mlas --release --lib # 1077 passed, 0 failed cargo test -p onnx-runtime-ep-cpu --release --lib # 1079 passed, 0 failed cargo clippy -p onnx-runtime-ep-cpu --features mlas --all-targets --release -- -D warnings # clean python3 scripts/check_platform_naming.py # PASS python3 scripts/check_dispatch_reachability.py # PASS python3 scripts/check_dispatch_manifest.py # PASS python3 scripts/check_feature_gate_coverage.py # PASS python3 .github/scripts/verify_documented_env_vars.py # PASS ``` --- ### Verified pre-existing CI failing set Re-checked against the unmodified baseline commit `0b872ed2f` ([CI run 31902589284](https://github.com/justinchuby/onnx-genai/actions/runs/31902589284)). Exactly these eight jobs fail on `main` with no changes applied: `CLI ORT (Linux x86_64)`, `CLI ORT (Windows x86_64)`, `CUDA compile (Linux x86_64)`, `CUDA compile (Windows x86_64)`, `Fast (Linux x86_64)`, `Rust (Windows ARM64)`, `Rust coverage (macOS arm64)`, `Rust quality`. This is the same set this PR shows — no job fails here that does not already fail on `main`. `Fast (Linux x86_64)` and `Rust quality` both fail in their `Check formatting` step: `cargo fmt --all -- --check` reports the identical 7 diffs on this branch and on unmodified `origin/main`, all in `onnx-genai-ort` and `onnx-runtime-ep-cuda`, none in a file this PR touches. ## CI baseline Rebuilt on top of `main` @ `400fbe246` (a plain `git merge origin/main`, no history rewrite). Unmodified `main` @ `400fbe246` fails exactly these 6 jobs ([run 31914831964](https://github.com/justinchuby/onnx-genai/actions/runs/31914831964)): | Job | Fails on unmodified `main` | |---|---| | `CLI ORT (Linux x86_64)` | yes | | `CLI ORT (Windows x86_64)` | yes | | `CUDA compile (Linux x86_64)` | yes | | `CUDA compile (Windows x86_64)` | yes | | `Rust (Windows ARM64)` | yes | | `Rust coverage (macOS arm64)` | yes | None are touched by this PR. `Fast (Linux x86_64)` and `Rust quality` previously failed on `main` too (a repo-wide `cargo fmt` drift, fixed on `main` by #1043); after merging current `main` into this branch both are green here, which confirms those earlier reds were never mine. The jobs this PR is actually accountable for -- `Fast (Linux x86_64)`, `Rust quality`, `EP conformance (Linux x86_64)`, `Rust coverage (Linux x86_64)`, `Miri unsafe-crate soundness`, `audit` and `codecov` -- are green. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Owner
Author
Requested changes — implemented in follow-up #1051#1027 was squash-merged into Host: RTX 4060 laptop, 20 logical CPUs, AVX2+FMA+F16C+AVX-VNNI, no AVX-512. Build qwen05b-symzp — on-disk
|
| 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 |
qwen14b-symzp — on-disk 8,549,241,669 B (≈7.96 GiB); memory measured on the branch build (accounting is deterministic)
| config | admitted | peak RSS | resident_f32_cache_bytes |
total_weight_bytes |
|---|---|---|---|---|
| QNBIT=0 borrowed | true | ~8.0 GB (≈ on-disk; WS observed to 7.95 GB) | 0 | 8,549,241,669 |
| QNBIT=1 MLAS | true | 25.1 GB | 6,995,312,640 | 15,544,554,309 |
resident_f32_cache_bytes goes 0 → 235 MiB (qwen05b) / 6.51 GiB (qwen14b) when the MLAS route is taken — the packed buffer is now visible to the ledger/governor where it read 0 in both modes before. qwen14b peak 25.1 GB reproduces the ~25.5 GB regression. The accounted number is MLAS's own sqnbit_packed_b_size (the budgetable persistent buffer); peak RSS runs higher because SQNBitPackedB also keeps a scales copy and packing transiently materialises a dense int4 copy (freed after). 14B decode ms/tok was not isolable on this contended host (a second worktree ran native inference concurrently); memory figures are per-process and clean.
Item 4 — can the packed buffer replace the mapped weights (→ ~1x)?
Measured: no, not without a loader redesign. The weights are a single shared read-only mmap of model.onnx.data owned centrally by WeightStore (kept alive via Arc); every initializer is a byte-range slice. Kernels get a borrowed TensorView per call with no per-initializer page-range info and no authority to unmap; other nodes on the borrowed fallback path plus embeddings/lm_head/layernorm still read the same file. Releasing pages would need the loader to expose per-initializer page ranges, a post-prepack hook, an OS-specific discard (DiscardVirtualMemory/madvise(MADV_DONTNEED)), and page-boundary care. Measurement confirms coexistence: 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. Worthwhile follow-up, out of scope for a kernel-local change.
Policy when the budget is too small
Admitted through the exact #987 gate: admit only when on_disk + packed ≤ residency ceiling. Otherwise 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 (byte-identical, only slower on x86). Demonstrated above with --host-ram-limit 500MB (admitted=false, peak held at the borrowed level).
Tests: cargo test -p onnx-runtime-ep-cpu --features mlas --lib matmul_nbits → 100 passed, 0 failed, 6 ignored (re-stated #979 footprint probe + new decline test).
This was referenced Aug 16, 2026
justinchuby
added a commit
that referenced
this pull request
Aug 16, 2026
…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
Owner
Author
The marginal value of this route dropped sharply once #1021 landed, and that changes the default questionWhen this PR was measured, the comparison was MLAS against a scalar borrowed int4 path, and the answer was 2.68x. #1021 then vectorised that path (AVX2/AVX-512). The number that matters for any "should
Marginal gain: 1.25x, not 2.68x. Text byte-identical. MLAS is also cheaper on fixed cost (20.2 vs 28.7 s CPU) -- the pack is repaid within prefill on this model -- so the win is real, just much smaller than it was. Against that, from #1051's accounting: admitting the route costs +603.6 MB on a 367 MB model (peak RSS 390 → 1002 MB, a 2.5x process footprint), and on Recommendation: keep
|
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
…nt copy) (#1104) ## Summary MLAS SQNBit CompFp32 beats our borrowed int4 decode path (#1021) by ~1.25x on CPU (#1027). This PR finds **where** the gap comes from, proves it is reachable **without** MLAS's session-lifetime resident packed copy, and ports it. **The 1.25x is register/N-blocking, not arithmetic and not layout.** Both kernels are pure f32 FMA (CompFp32 uses no VNNI), so arithmetic is equal. Reading the MLAS AVX2 M==1 kernel (`sqnbitgemm_kernel_avx2.cpp`) against our borrowed path (`matmul_nbits.rs`), the concrete differences are: | | MLAS SQNBit M==1 | our borrowed path (before) | | --- | --- | --- | | output columns per pass | **4** (`NCols4`), 4 live accumulators | 1, single accumulator | | activation load | loaded once, **reused across 4 columns** | reloaded per column | | horizontal reduction | **once per column** at the end | **once per block** | | nibble unpack | shuffle-free (prepacked `\|v0 v16\|…`) | `unpacklo/unpackhi` interleave | | resident cost | session-lifetime packed copy (~2x int4 bytes) | **none** (zero-copy mmap) | The first three are pure inner-loop restructuring that need **no** repack. Only the shuffle-free unpack is tied to MLAS's prepacked layout. ## What this implements `borrowed_affine_int4_matmul_nblock`: reads the **same zero-copy mmap int4 layout** (no repack, no resident copy) but processes up to four output columns per pass with four independent f32 accumulators, loads each activation vector once and reuses it across the group, folds each block's scale into the running accumulator with a single FMA, carries the zero-point affine correction in scalar, and does one horizontal reduction per column. The nibble unpack is byte-for-byte identical to the per-column AVX2 helper. Selected for decode on AVX2-capable hosts behind the same-binary A/B toggle **`ONNX_GENAI_CPU_MM_INT4_NBLK`** (default off; the toggle is a read-only env probe, production is the only writer — no new process-global mutable state). ## Where the 1.25x comes from — evidence All measured on this host (RTX 4060 laptop, 20 logical CPUs, AVX2+FMA+AVX-VNNI, **no AVX-512, CPU only**), `--backend native`, `ONNX_GENAI_CPU_DECODE_THREADS=8`, decode isolated as **88-tok − 8-tok**, **process CPU time**, **min of 3**, on `models\qwen05b-symzp`: | path | short 8 tok | long 88 tok | decode (80 tok) | **CPU s/token** | **peak RSS** | | --- | --- | --- | --- | --- | --- | | ours (borrowed int4) | 6.7 s | 55.5 s | 48.8 s | **0.61** | **389 MB** | | **nblk (this PR)** | 4.4 s | 37.9 s | 33.5 s | **0.42** | **389 MB** | | MLAS SQNBit CompFp32 | 5.6 s | 36.2 s | 30.6 s | **0.38** | **1059 MB** | **How much is reachable without a resident copy:** nblk recovers **~83%** of MLAS's decode advantage — `(0.61−0.42)/(0.61−0.38)` — at the borrowed path's **389 MB**, versus MLAS's **1059 MB**. The residual nblk↔mlas gap (0.42 vs 0.38) is **within nblk's own run-to-run spread** (long runs 37.9 / 43.7 / 45.8 s), i.e. below the measurement floor. Because nblk reads the **normal mmap layout** and still captures essentially all of the win, **the advantage is inner-loop register blocking, not layout**; the shuffle-free-unpack contribution is at or under the noise floor on this host and does not justify a transient repack. ## The 14B — the model that needed it most On `models\qwen14b-symzp` MLAS's resident cache is **declined** at the default ceiling: predicted `resident_f32_cache_bytes = 17,925,488,640` (16.7 GB) exceeds the residency ceiling, `f32_weight_cache_admitted=false`, so the MLAS arm peaks at **8189 MB** — identical to borrowed. **MLAS gives the 14B nothing.** nblk needs no admission (it is not a resident cache), so it applies here: | path | decode (40-8 tok) | **CPU s/token** | **peak RSS** | | --- | --- | --- | --- | | ours (borrowed int4) | 442.5 s | **13.83** | 8220 MB | | **nblk (this PR)** | 304.1 s | **9.50** | 8222 MB | **1.46× faster decode at the same 8.2 GB**, byte-identical — exactly the model MLAS could not help. ## Byte-identity Generated text is byte-identical (SHA-256 of `--raw` output) across **borrowed / nblk / mlas** on: - `qwen05b-symzp` (symmetric) — `12B539F850DFE618` - `qwen05b-q4-zp` (asymmetric) — `3EA4EABE60BB60D1` - `qwen14b-symzp` (40-tok) — `3C517757A47C1395` ## Gates - `cargo test -p onnx-runtime-ep-cpu --lib`, five consecutive runs: **1308 / 1308 / 1308 / 1308 / 1308 passed, 0 failed, 11 ignored** each. - `cargo clippy -p onnx-runtime-ep-cpu --lib -- -D warnings`: clean. - New parity test `nblock_matches_per_column_borrowed_path` (symmetric + asymmetric, single/multi-row, non-multiple-of-4 tail group). `cargo test --features mlas --lib matmul_nbits`: 111 passed, 0 failed, 6 ignored. ## Notes - Default off; this is an opt-in A/B path, like #1021 and #1027. Flipping it on by default is a follow-up once it has soaked, but the measurements above show it is a strict win (faster decode, identical footprint, byte-identical output). - Not absorbed: the shuffle-free in-block nibble order. It requires either MLAS's resident prepack (refused) or a transient repack, and its isolated contribution is under the measurement floor here — a well-measured "not worth a repack on this host." Refs #1021, #1027, #1051, #1056 Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
## Share one MLAS SQNBit packed buffer per weight (#1056) Refs #1027 (MLAS SQNBit route), #1051 (packed-buffer accounting), #1056 (this dedup). ### The problem `MatMulNBits` int4 `accuracy_level=0` nodes route to MLAS SQNBit CompFp32 (#1027). The executor's `KernelCache` is **shape-keyed**, so each node compiles two kernel instances — prefill (`m > 1`) and decode (`m == 1`) — and before this change **each instance packed its own full copy of the same constant weight**. The resident packed footprint was therefore `2x` the 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 to reach a weight packs it once; the sibling instance takes the same `Arc`. The session now holds **one** packed copy per weight. * Keyed on the mmap **address** plus every pack-determining shape/param, so a same-address different-shape weight (allocator recycling a freed address — the #845/#1079 hazard) misses rather than serving the wrong bytes. * `clear_mlas_packed_caches()` runs on `Executor` drop — the **same lifetime boundary** as `weight_transpose::clear_all` — closing the same-address/same-shape/across-lifetimes window. * Accounting is updated **in the same commit**: `MLAS_PACKED_DECODE_INSTANTIATIONS` goes `2 -> 1`, so `resident_dequant_f32_cache_bytes` (the plan's prediction) equals `mlas_sqnbit_packed_live_bytes` (the actual allocation). The existing accounting test additionally asserts **pointer identity** of the shared `Arc` across the prefill and decode instances. * **No new process-global mutable state that tests mutate** (#983/#1033/#1079): production reads only the global store; tests use a `cfg(test)` **thread-local** store (each libtest thread gets a private 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) vs `live_total` (profiler actual `SQNBIT_PACKED_LIVE_BYTES`), measured on-model: | model / arm | predicted (bytes) | actual live (bytes) | ratio | |---|---:|---:|---:| | qwen05b after, admitted, multi-token | 316,443,904 | 316,443,904 | **1.00** | | qwen05b before, admitted, multi-token | 632,887,808 | 632,887,808 | 1.00 | | qwen14b after, admitted (18GiB ceiling) | 8,962,744,320 | 8,962,744,320 | **1.00** | The dedup test (`int4_acc0_mlas_packed_accounting_equals_actual_allocated`) that ties the plan's prediction 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=` on `qwen05b-symzp` (169 weight boundaries): | run | before (origin/main) | after (this branch) | |---|---:|---:| | 1-token (single activation shape) | 169 | 169 | | multi-token (prefill + decode shapes) | **338** | **169** | 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,808` for it — over-report, the safe direction; after, both the pack count and 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 `PeakWorkingSet64` while running; CPU time = `TotalProcessorTime`. **qwen05b-symzp** (weights 366,846,066 B), multi-token autoregressive run: | arm | packs | accounted | live | peak RSS | admitted | |---|---:|---:|---:|---:|:--:| | route OFF (`QNBIT=0`) | – | 0 | – | 485.9 MB | – | | route ON **before** | 338 | 632,887,808 | 632,887,808 | 1135.9 MB | yes | | route ON **after** | 169 | 316,443,904 | 316,443,904 | **818.5 MB** | yes | 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`): | arm | accounted (predicted) | live | expanded footprint | peak RSS | verdict | |---|---:|---:|---:|---:|:--:| | **before**, default ceiling | 17,925,488,640 | 0 (declined) | 26,474,730,309 | 8643.4 MB | **declined** | | **after**, default ceiling | 8,962,744,320 | 0 (declined) | 17,511,985,989 | 8669.3 MB | **declined** | | **after**, ceiling 18 GiB | 8,962,744,320 | 8,962,744,320 | 17,511,985,989 | 17,606.9 MB | **admitted** (ratio 1.00) | **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 RAM` ceiling of 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 peak 17,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: | prompt | before | after (declined) | after (admitted) | |---|---|---|---| | qwen05b, "…relativity…" 48 tok | `EF7CA14F…` | `EF7CA14F…` | – | | qwen05b, raw "a" 8 tok | `FD4972FF…` | `FD4972FF…` | – | | qwen14b, "…relativity…" 16 tok | `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 `LazyLock` global; tests use a `cfg(test)` thread-local store restored automatically by each libtest 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, and the `NO_SHARD` A/B) goes through the shared store. The only non-shared fallback is a weight with no stable 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`._ Co-authored-by: justinchuby <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
…ed sums) Replace the row-serial borrowed int4 prefill path's two per-token overheads with a structural rewrite behind a default-off A/B toggle (ONNX_GENAI_CPU_MM_INT4_PREFILL), per the method used in #1021/#1027/#1104/#1116: - One fork-join over disjoint column strips for the whole prefill, instead of the row-serial path's m per-row fork-joins. - The per-block �ctivation_sums Vec is hoisted out of the per-row loop to a single m * block_count allocation, independent of the weight size. Rows are still visited outer-most, so a column's packed bytes are re-read once per row: weight traffic and the per-element k-reduction order are unchanged, so output is byte-identical to the row-serial path (both call the shared �orrowed_int4_output_element). No resident buffer is added (peak RSS unchanged), satisfying the #1056/#1117 no-session-buffer constraint. GEMM blocking (reusing a column's bytes across a tile of rows) is deliberately NOT included here: measured within run-to-run noise on qwen05b and with no signal on qwen14b (5-rep interleaved), so it is left as an unproven follow-up to #1117 that can be added or dropped cleanly. Measured win is model-size dependent: qwen05b prefill slope 1.219 -> 0.842 CPU s/token (median of 5, ~1.45x); qwen14b shows no measurable change (medians 8.69/8.62 off/on, within a 30-40% within-arm spread) because the fixed per-row overheads this removes are a negligible fraction of the 14B's larger per-row compute. Refs #1117 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: d60eb808-7cc6-4abc-b48d-2a6dd3841624
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Kernel::executeevaluates thebits == 4 && accuracy_level == 0borrowed zero-copy branch beforetry_mlas_sqnbit. That branch order made theaccuracy_level = 0→ MLAS SQNBit CompFp32 route unreachable dead code — the very routetry_mlas_sqnbit's own doc comment describes as intended:On x86_64 the consequence is severe:
borrowed_affine_int4_matmul's vectorized block dot is#[cfg(target_arch = "aarch64")], so on x86 everybits=4, accuracy_level=0node ran a scalar nibble-unpack GEMV, and form > 1it re-streamed the whole weight once per activation row with no GEMM blocking.accuracy_level = 0is exactly what Foundrycuda-gpuint4 exports emit, so this is the common case, not a corner.Reproduction and measurement method
Single-node
com.microsoft::MatMulNBitsONNX models (constant B/scales/zero-points) executed through both our native CPU EP and the real ORT 1.27 CPU EP in one process, interleaved, viabench_generic.--ort-intra-threads 8againstONNX_GENAI_CPU_DECODE_THREADS=8.ONNX_GENAI_CPU_MM_MLAS_QNBIT=0/1, so compiler, ISA, allocator and host load are identical across the pair.Results
bits = 8rows are untouched by the gate (bits == 4) and measured flat, as expected.Against ORT the same node moves from 12x–78x slower to 1.5x–4.5x slower. This PR does not claim parity with ORT — see "What this PR does not fix".
Correctness
Parity vs ORT becomes bit-identical, because both runtimes now execute the same MLAS kernels:
Setup cost is reported, not hidden
The one-time MLAS shard pack for the K=3584 N=3584 block-32 weight, from
ONNX_GENAI_PROFILE_MM=1:55.8 ms once for 6.4 MB, against ~4.1 ms saved per call → repaid in ~14 calls, then amortized for the session. For decode workloads that is a fraction of one token.
Design: an explicit ownership predicate, not a branch reorder
mlas_sqnbit_owns_fp32_compute(can_prepack, has_zero_points)states the rule where it can be read and tested:can_prepack— B/scales/zero-points are graph constants — so the packed buffer is built once per session. With dynamic weights MLAS would repack on every call with no cache to amortize against, so those explicitly stay on the borrowed zero-copy path. This is the deliberate decline half of the policy, covered bymatmulnbits_int4_acc0_dynamic_weight_keeps_borrowed_path.sqnbit_packed_b_sizeis MLAS's own "do I have a kernel for this shape on this host" probe. When MLAS says no (e.g.block_size = 8), the borrowed path remains the fallback — nothing is left stranded.mlasfeature the predicate isfalse, so non-MLAS builds are byte-for-byte unchanged.#979 is preserved
#979 removed the resident f32
weight_nkexpansion (~8x the file size in RAM). MLAS's packed buffer is int4-sized, so that expansion is still never built. The #979 regression tests now assert the memory invariant directly (weight_nkstays empty) alongside the route actually taken, where the expected route is derived from the same predicate production code uses — so on hosts/shapes where MLAS declines they stay exact assertions rather than degrading into tautologies. Under--features mlasthe route proof reads per-kernel state (mlas_shards/mlas_packed), not the process-global test counters, which are shared with tests running in parallel and cannot support a negative assertion.What this PR does not fix (measured, reported, not hidden)
accuracy_level = 0, all shapesaccuracy_level = 4, M=1sqnbit_decode_min()short-circuits M=1 onto the hand int8 kernel, which needs AVX-VNNI to be fast. Follow-up PR.accuracy_level = 4, M=1, asymmetric, AVX2-onlySQ4BitGemmM1Kernel_CompInt8_avx2asymmetric kernel is numerically broken (already guarded in-tree). Correct decline; documented gap.bits = 8accuracy_level = 4mlas) buildsmlasfeature is not enabled for published wheels/CLI — see below.The
mlasfeature gap is the biggest remaining exposureThis fix only helps builds compiled with
--features mlas, which is opt-in and not enabled for shipped artifacts. On a default x86 build the borrowed scalar GEMV is still what runs. Two follow-ups are needed and are being handled separately: (1) an AVX2 int4 CompFp32 GEMV/GEMM so default builds are not scalar, and (2) a decision on enablingmlasby default forx86_64-unknown-linux-gnuartifacts.Validation
Both
cfg(feature = "mlas")andcfg(not(feature = "mlas"))variants of the predicate exist, so the feature-gate coverage rule is satisfied.Review
Independent Rubber Duck review (Opus, read-only): APPROVE.
One MINOR finding:
Int4Acc0RouteProbe::assert_fast_routederived its expectation frommlas_sqnbit_owns_fp32_compute— the very predicate under test — so it caught "the code disagrees with its own policy" but not "the policy regressed to never choosing MLAS".Fixed in
4c599eafd:int4_acc0_constant_weight_reaches_mlas_on_x86_64pins the concrete case this PR exists for with a hardcoded expectation — on x86_64 with the vendored MLAS, a constant-weight symmetric int4accuracy_level = 0node at block size 32 must reach MLAS SQNBit. Reintroducing the branch order that made that route dead code now fails a test even if the predicate is edited to agree with it.CI
mainis currently red for reasons that predate and are untouched by this PR. Verified by checking out unmodifiedorigin/mainin this worktree:cargo fmt --all -- --checkflagscrates/onnx-genai-ort/src/lib.rsandcrates/onnx-runtime-ep-cuda/src/kernels/matmul_nbits.rsonorigin/mainitself. This PR changes exactly one file,crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs, and the flagged set is byte-identical before and after.CLI ORT (Linux/Windows)→ Build onnx-genai-cli,CUDA compile (Linux)→ Verify CUDA test inventory,CUDA compile (Windows)→ Clippy CUDA EP,Rust (Windows ARM64)→ Test cross-platform offline crates. The last 8CIruns onmainall concludefailure.Locally green on this branch:
Verified pre-existing CI failing set
Re-checked against the unmodified baseline commit
0b872ed2f(CI run 31902589284).
Exactly these eight jobs fail on
mainwith no changes applied:CLI ORT (Linux x86_64),CLI ORT (Windows x86_64),CUDA compile (Linux x86_64),CUDA compile (Windows x86_64),Fast (Linux x86_64),Rust (Windows ARM64),Rust coverage (macOS arm64),Rust quality.This is the same set this PR shows — no job fails here that does not already
fail on
main.Fast (Linux x86_64)andRust qualityboth fail in theirCheck formattingstep:cargo fmt --all -- --checkreports the identical 7diffs on this branch and on unmodified
origin/main, all inonnx-genai-ortand
onnx-runtime-ep-cuda, none in a file this PR touches.Re-validated against merged
main(eac0ee24e)The other seven PRs in this series are now merged. I rebuilt both arms from source on the
same host and confirmed the binaries actually differ (
md5 a61b3130formain,8e842c67with this PR) before trusting the numbers, then measuredbench_generic --runs 40 --warmups 10 --native-threads 8 --ort-intra-threads 8:accuracy_level=0, b4, bs32, sym)mainvs ORTaccuracy_level=0is the only remaining int4 regression onmain--acc4is already1.92x (M=1) / 2.38x (M=128) after #1028. Numeric parity is
PASS(bit-identical) in everyrun. Absolute timings on this shared box drift 1.5-2x between runs, so the ratios matter,
not the milliseconds.
CI baseline
Rebuilt on top of
main@400fbe246(a plaingit merge origin/main, no history rewrite).Unmodified
main@400fbe246fails exactly these 6 jobs(run 31914831964):
mainCLI ORT (Linux x86_64)CLI ORT (Windows x86_64)CUDA compile (Linux x86_64)CUDA compile (Windows x86_64)Rust (Windows ARM64)Rust coverage (macOS arm64)None are touched by this PR.
Fast (Linux x86_64)andRust qualitypreviously failed onmaintoo (a repo-widecargo fmtdrift, fixed onmainby #1043); after merging currentmaininto this branch both are green here, which confirms those earlier reds were never mine.The jobs this PR is actually accountable for --
Fast (Linux x86_64),Rust quality,EP conformance (Linux x86_64),Rust coverage (Linux x86_64),Miri unsafe-crate soundness,auditandcodecov-- are green.Ratio convention (added post-merge for clarity)
Columns
before/ORTandafter/ORTare ours/ORT: >1 means we are slower, soafter/ORT 1.54is 1.54x slower than ORT, not faster.speed-upis ours-before/ours-after (this PR's own gain) and never involves ORT. p50 of 9 runs after 3 warmups, interleaved A/B, 8 threads on both sides, steady state (weights already packed; packing is not amortised into these numbers). Shared, contended host: treat <1.2x as noise.