Skip to content

perf(ep): tell kernels which of their inputs ORT holds constant - #1110

Merged
justinchuby merged 3 commits into
mainfrom
squad/roy-plugin-constant-weights
Aug 17, 2026
Merged

justinchuby merged 3 commits into
mainfrom
squad/roy-plugin-constant-weights

Conversation

@justinchuby

@justinchuby justinchuby commented Aug 17, 2026 •

Copy link
Copy Markdown
Owner

The bug

Kernel::set_constant_inputs is how a kernel learns that an input is a weight ORT owns for the life of the session. Kernels use it to decide whether a prepack may be built once and kept, and MatMulNBits goes further: mlas_sqnbit_owns_fp32_compute(can_prepack, …) gates the MLAS SQNBit path itself on the same flag.

It was called from exactly one place in the workspace — onnx-runtime-session's executor. The ORT plugin EP never called it. So on the only path a real ORT model takes, every kernel saw every input as a runtime tensor: MatMulNBits rebuilt its packed weights on every Run and ran the slower non-SQNBit kernel, and QLinearMatMul re-packed B per call. A decoder paid a load-time cost per token.

Why the existing test could not be reused

node_inputs_all_routable already distinguishes weights, as "no producer and not a graph input". That is right for the whole model at GetCapability time and wrong at Compile time: ORT hands a fused node's subgraph over with the initializers it kept inside listed as graph inputs of that subgraph, because they are inputs of the fused node. Instrumenting the compile path on a MatMulNBits model prints

[const] op=MatMulNBits flags=[false, false, false] is_graph_input=[true, true, true] has_producer=[false, false, false]

Graph_GetInitializers still tells the two apart, so OutboundGraphReader now records initializer names (no tensor data: the existing read_initializers_int64 deliberately copies only small int64 tensors, and a 1 GB B must not be copied to answer a yes/no question), and constant_input_flags keys against that set.

Measurement

Interleaved A/B against plain ORT through the ORT session API — the same generated single-node model loaded twice, once with this EP registered and session.disable_cpu_ep_fallback=1, once with no EP appended at all. 3 warmups, then 41 interleaved iterations each side, p50/p90 of Run only (input OrtValues built once). Assignment is asserted before any timing, so a ratio is never reported for a node ORT actually ran.

Host: AMD EPYC 9V74, 32 vCPU / 16 cores, AVX2+FMA+F16C (no AVX-512/VNNI), ORT 1.27.0, release build with --features onnx-runtime-ep-cpu/mlas, K=N=2048. The box is shared, so p90s are contended; the ours ms columns are the load-bearing evidence and the ORT column moves between runs.

case ours ms before ours ms after speedup ours/ORT before ours/ORT after
MatMulNBits int4 M=1 1.144 0.425 2.7x 15.0 2.34
MatMulNBits int4 M=128 108.4 7.79 13.9x 91.6 6.36
MatMulNBits int4 f16-act M=1 3.659 0.416 8.8x 16.0 1.80
MatMulNBits int8 M=1 8.226 4.845 1.7x 3.85 2.53
MatMulNBits int8 M=256 47.98 12.04 4.0x 5.03 1.05
QLinearMatMul u8 M=1 2.234 0.079 28x 43.6 1.84
QLinearMatMul u8 M=128 13.32 11.11 1.2x 2.15 3.78
QLinearMatMul i8 M=1 1.968 0.091 21.6x 0.52 0.055
MatMul f32 M=1 (control) 0.147 0.178 — 1.72 1.49
MatMul f32 M=128 (control) 9.675 8.698 — 1.36 1.43
MatMul f16 M=1 (control) 1.368 2.358 — 1.08 0.37
MatMul f16 M=128 (control) 21.04 22.64 — 7.04 2.84

The four dense cases are the control: they declare both operands as graph inputs, so they have no initializer, no flag changes for them, and they move only with host noise. Every case that moved has a constant weight.

Cold session-creation time is reported by the same harness (cold_ours_ms) and does not regress: the prepack moved from per-Run to first-Run, not to CreateSession.

Still losing after this change

Reporting all of it, per the standing rule:

  • MatMulNBits int4 M=128 at 6.36x and M=1 at 2.34x.
  • QLinearMatMul u8 M=128 at 3.78x (u8 M=1 is now 1.84x).
  • MatMul f16 M=128 at 2.84x — no initializer involved, so untouched here.
  • MatMul f32 at 1.43-1.49x — likewise.

These are kernel gaps, not wiring gaps, and they stay open on my task list.

Separate finding, not fixed here: the plugin cdylib is built without the mlas feature by default, and docs/architecture/CROSS_PLATFORM.md documents that feature as x86-64-Linux-only. Rebuilding the same benchmark against the default cdylib gives, after this fix, MatMulNBits int4 M=128 at 80x and QLinearMatMul u8 M=128 at 56x ours/ORT: the pure-Rust paths are the ones a portable build actually ships, and they are far behind. That is a distinct piece of work and needs its own PR.

Tests

  • constant_weights_are_reported_to_kernels_as_constant (new, plugin_ort_e2e) — reads the counter this PR exports from the very cdylib ORT loaded and asserts symmetric int4 reports 2 constant inputs, asymmetric int8 reports 3, and the all-graph-input dense case reports 0. A wiring that marks everything constant fails the third assertion; caching a prepack of an activation is wrong, not merely slow. Falsified by reverting constant_input_flags to the producer/is_graph_input form: reported 0 constant inputs, expected 2.
  • initializers_are_constant_even_when_the_subgraph_calls_them_inputs and nothing_is_constant_without_an_initializer_list (new, onnx-runtime-ep-plugin unit) — hardware-independent, pin the exact shape ORT presents and the absent-optional-input case.
  • no_matmul_family_node_escapes_to_the_ort_cpu_ep (extended) — every case now runs a second time in the same session with the activation rotated by one element, and is compared against ORT again. A weight cache still matches ORT; anything activation-derived that outlived the call does not. ORT is the oracle for whether a case can detect staleness at all (a saturating u8 output cannot), and the suite asserts at least 5 of the 10 cases are activation-sensitive so the check cannot go vacuous.
  • plugin_path_ab_vs_plain_ort (new, #[ignore] + NXRT_MM_BENCH=1) — the harness that produced the table, committed so the numbers are reproducible rather than asserted.
  • plugin_export_abi — the two new exports are added to both the required-symbol list and the unexpected-symbol filter.

Verification

  • cargo test -p onnx-runtime-ep-cpu-plugin — 51 + 9 + 6 + 1 passed, 1 ignored (the benchmark)
  • cargo test -p onnx-runtime-ep-plugin --lib — 231 passed (was 229)
  • cargo fmt --all -- --check clean (given style: reformat the int4 tail-block slice #1104 left unformatted #1109, which repairs main)
  • cargo clippy --all-targets clean with and without --features onnx-runtime-ep-cpu/mlas

Post-review fix (commit 2): overridable initializers are not constant

Independent review found a correctness hole in commit 1, and it was right.

Graph_GetInitializers is documented in ORT's own header as including "constant and non-constant initializers". From ONNX IR version 4, an initializer whose name also appears in the graph's input list is only a default: the caller may hand a different tensor in on any Run. ORT says so at load time:

[W:onnxruntime:, graph.cc:1419 Graph] Initializer B appears in graph inputs and will not be
treated as constant value/weight.

Keying the flags on that name list alone therefore marked such a value constant, and MatMulNBits would have cached a session-lifetime prepack of the default weight and returned its answer for every later Run — a wrong result, not a slow one.

Fix: each name from Graph_GetInitializers is now filtered through ValueInfo_IsConstantInitializer (ORT ≥1.23, reads no tensor data). Fail-closed: if the entry point is missing or the call errors, the initializer is treated as non-constant. A false negative costs a repeated prepack; a false positive costs correctness.

New regression test an_overridable_initializer_is_not_treated_as_a_constant_weight builds exactly that model — B and scales declared both as initializers and as graph inputs — and asserts two independent things:

  1. the EP reports 0 constant inputs (the classification), and
  2. running the session twice with two different B payloads matches plain ORT run for run, with the test first proving via ORT that the two payloads do produce different outputs (the consequence — a stale prepack would return run one's answer twice).

Falsified by deleting the is_constant filter and rebuilding the cdylib:

assertion `left == right` failed: an initializer that is also a graph input may be replaced
on any Run, but this EP reported 2 of them as constant weights

The two doc comments that asserted "every initializer is constant" are corrected, and the unit test is renamed to constant_initializers_are_flagged_even_when_the_subgraph_calls_them_inputs.

The measured win survives the stricter gate

Re-measured after the fix on the same host, same build (--features onnx-runtime-ep-cpu/mlas, release, K=N=2048, 3 warmups + 41 interleaved iterations, p50 ms):

case ours before PR ours after commit 1 ours after commit 2
nbits4_m1 1.144 0.425 0.400
nbits4_m128 108.4 7.79 8.80
nbits4_f16_m1 3.659 0.416 0.437
nbits8_m256 47.98 12.04 13.86
qlinear_u8_m1 2.234 0.079 0.092
qlinear_i8_m1 1.968 0.091 0.092

The benchmark models declare their weights as initializers only, so they are still constant under the stricter rule; the residual movement is host noise on a shared machine (ORT's own p50 for nbits4_m1 moved 0.125 → 0.077 ms between the two runs, so the ratio column is the noisier of the two and the absolute ours column is the honest comparison here).

Verification (re-run)

  • cargo test -p onnx-runtime-ep-cpu-plugin --test plugin_ort_e2e — 52 passed, 1 ignored (was 51 + 1)
  • cargo test -p onnx-runtime-ep-plugin --lib — 231 passed
  • cargo fmt --all -- --check clean (rebased onto style: reformat the int4 tail-block slice #1104 left unformatted #1109, now merged)
  • cargo clippy -p onnx-runtime-ep-cpu-plugin -p onnx-runtime-ep-plugin --all-targets clean with and without --features onnx-runtime-ep-cpu/mlas

@codecov

codecov Bot commented Aug 17, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.90909% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.86%. Comparing base (805890e) to head (cc873b5).

Files with missing lines Patch % Lines
crates/onnx-runtime-ep-plugin/src/graph_reader.rs 78.18% 5 Missing and 7 partials ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1110      +/-   ##
==========================================
- Coverage   80.40%   79.86%   -0.55%     
==========================================
  Files         369      367       -2     
  Lines      160217   157553    -2664     
  Branches   160217   157553    -2664     
==========================================
- Hits       128827   125823    -3004     
- Misses      26656    27007     +351     
+ Partials     4734     4723      -11     
Flag Coverage Δ
cli-ort-linux 83.79% <ø> (ø)
cli-ort-windows 83.40% <ø> (+0.09%) ⬆️
mlas ?
offline 79.71% <90.90%> (-0.49%) ⬇️

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

Files with missing lines Coverage Δ
crates/onnx-runtime-ep-cpu-plugin/src/lib.rs 76.19% <100.00%> (+1.83%) ⬆️
crates/onnx-runtime-ep-plugin/src/ep.rs 83.39% <100.00%> (+0.89%) ⬆️
crates/onnx-runtime-ep-plugin/src/graph_reader.rs 69.63% <78.18%> (+0.70%) ⬆️

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

@github-actions

github-actions Bot commented Aug 17, 2026 •

Copy link
Copy Markdown

✅ Benchmarks — No Regression

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
✅ matmul/large_generic_bf16_threads=8/32x1024x1024 2.13 ms 2.40 ms +12.7%
✅ reduce_mean/small_f32_threads=1-internal/4096 18.84 µs 19.08 µs +1.3%
✅ sampling_latency/min_p_per_token 206.94 µs 203.63 µs -1.6%
✅ gather/large_bf16_threads=1-internal/131072 15.17 µs 14.51 µs -4.3%
✅ tokenization/encode_tokens_per_second 419.55 µs 400.49 µs -4.5%
✅ kv_cache/alloc_dealloc_pages 39.07 µs 37.00 µs -5.3%
✅ logit_processing/seven_processor_chain_per_step 315.36 µs 297.67 µs -5.6%
✅ qwen3_sampling_processors/top_k_top_p_fast 663.88 µs 626.62 µs -5.6%
✅ qwen3_sampling_processors/top_p_fast_after_top_k 516.60 µs 485.80 µs -6.0%
✅ grammar_masking/llguidance_compute_mask/32 87.06 µs 81.80 µs -6.0%
✅ sampling_latency/top_p_per_token 391.09 µs 363.77 µs -7.0%
✅ qwen3_sampling_processors/top_k_full_sort_baseline 2.14 ms 1.98 ms -7.5%
✅ reduce_mean/large_f32_threads=1-internal/262144 1.13 ms 1.04 ms -8.1%
✅ block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1 176.28 µs 161.93 µs -8.1%
✅ add/large_bf16_threads=1-internal/4194304 2.45 ms 2.24 ms -8.7%
✅ qwen3_sampling_processors/top_k_top_p_full_sort_baseline 5.91 ms 5.32 ms -10.0%
✅ qwen3_sampling_processors/top_k_partial_selection 152.14 µs 136.85 µs -10.0%
✅ matmul/large_generic_bf16_threads=1/32x1024x1024 2.19 ms 1.94 ms -11.4%
✅ matmul/large_generic_f32_threads=1/32x1024x1024 9.84 ms 8.68 ms -11.8%
✅ qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline 3.74 ms 3.28 ms -12.2%
✅ matmul/medium_generic_f32_threads=1/32x512x512 2.46 ms 2.15 ms -12.6%
✅ tokenization/decode_tokens_per_second 8.18 ms 7.10 ms -13.2%
✅ add/small_f16_threads=1-internal/1024 707.3 ns 605.8 ns -14.3%
🟢 add/medium_bf16_threads=1-internal/262144 124.38 µs 104.81 µs -15.7%
🟢 matmul/small_generic_f32_threads=1/1x256x256 42.59 µs 35.88 µs -15.8%
🟢 reduce_mean/medium_f32_threads=1-internal/65536 297.57 µs 249.11 µs -16.3%
🟢 gather/small_f16_threads=1-internal/4096 617.7 ns 515.3 ns -16.6%
🟢 gather/medium_f32_threads=1-internal/32768 4.80 µs 3.95 µs -17.6%
🟢 matmul/medium_generic_f32_threads=8/32x512x512 1.13 ms 933.05 µs -17.7%
🟢 add/small_f32_threads=1-internal/1024 261.4 ns 212.3 ns -18.8%
🟢 gather/small_f32_threads=1-internal/4096 856.1 ns 686.9 ns -19.8%
🟢 matmul/small_generic_bf16_threads=1/1x256x256 36.76 µs 29.08 µs -20.9%
🟢 matmul/large_generic_f16_threads=8/32x1024x1024 100.59 µs 79.39 µs -21.1%
🟢 matmul/large_generic_f16_threads=1/32x1024x1024 95.17 µs 74.09 µs -22.2%
🟢 gather/small_bf16_threads=1-internal/4096 629.4 ns 484.0 ns -23.1%
🟢 add/large_f16_threads=1-internal/4194304 2.63 ms 2.02 ms -23.1%
🟢 sampling_latency/greedy_per_token 4.20 µs 3.16 µs -24.8%
🟢 sampling_latency/top_k_per_token 67.63 µs 50.83 µs -24.8%
🟢 gather/large_f32_threads=1-internal/131072 53.80 µs 39.58 µs -26.4%
🟢 matmul/small_generic_f16_threads=1/1x256x256 39.43 µs 28.60 µs -27.5%
🟢 add/small_bf16_threads=1-internal/1024 685.4 ns 493.1 ns -28.1%
🟢 block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1 654.51 µs 466.37 µs -28.7%
🟢 gather/medium_bf16_threads=1-internal/32768 3.51 µs 2.44 µs -30.5%
🟢 matmul/small_generic_f16_threads=8/1x256x256 40.67 µs 28.12 µs -30.9%
🟢 add/large_f32_threads=1-internal/4194304 905.59 µs 616.52 µs -31.9%
🟢 gather/medium_f16_threads=1-internal/32768 3.80 µs 2.49 µs -34.6%
🟢 matmul/small_generic_bf16_threads=8/1x256x256 45.42 µs 29.26 µs -35.6%
🟢 matmul/large_generic_f32_threads=8/32x1024x1024 5.82 ms 3.63 ms -37.6%
🟢 matmul/medium_generic_f16_threads=1/32x512x512 45.85 µs 27.95 µs -39.1%
🟢 matmul/medium_generic_bf16_threads=1/32x512x512 827.43 µs 492.56 µs -40.5%
🟢 matmul/medium_generic_bf16_threads=8/32x512x512 608.27 µs 355.77 µs -41.5%
🟢 add/medium_f16_threads=1-internal/262144 190.74 µs 106.91 µs -43.9%
🟢 matmul/medium_generic_f16_threads=8/32x512x512 51.89 µs 28.54 µs -45.0%
🟢 gather/large_f16_threads=1-internal/131072 21.84 µs 11.72 µs -46.4%
🟢 block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024 97.06 µs 48.76 µs -49.8%
🟢 add/medium_f32_threads=1-internal/262144 53.97 µs 25.39 µs -53.0%
🟢 block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024 1.19 ms 522.87 µs -56.1%
🟢 matmul/small_generic_f32_threads=8/1x256x256 81.29 µs 32.80 µs -59.7%
🟢 block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024 128.47 µs 42.88 µs -66.6%

Visual flags: ⚠️ ≥ 15% slower, 🔴 ≥ 30% slower — calibrated against measured runner noise (~27% worst-case on multi-threaded matmul)

Host info
CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 3.33 3.38 4.96 }
What this cannot catch
  • Regressions in code paths not covered by these benchmarks (e.g., end-to-end decode with a real model)
  • Sub-threshold regressions that compound over multiple PRs
  • Performance changes that only manifest under GPU execution
  • Latency changes in the ORT integration path (these benchmarks exercise the native Rust kernels)

@justinchuby justinchuby left a comment

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rubber-duck review — falsification attempt on PR #1110

Verdict: Request changes. One major soundness bug survived (and was empirically reproduced): the fix marks ONNX IR≥4 non-constant / overridable initializers as session constants, which ORT explicitly forbids. Everything else in the PR held up under attack — the mechanism is real, the wiring is correct, the tests are genuine falsifiers, and the perf numbers reproduce.

I reviewed from a detached worktree at 41627b6e1, built the cdylib with --features onnx-runtime-ep-cpu/mlas, and ran the suites against it (NXRT_CPU_PLUGIN_PATH pointed at a uniquely-named copy so I never touched the author's artifacts).


MAJOR — Claim 1 is false: Graph_GetInitializers ≠ "constant for the session lifetime"

read_initializer_names (crates/onnx-runtime-ep-plugin/src/graph_reader.rs:776) records every name returned by Graph_GetInitializers, and constant_input_flags (crates/onnx-runtime-ep-plugin/src/ep.rs:976) marks any node input in that set as constant. But Graph_GetInitializers does not return only constants. From ORT's own header (onnxruntime_c_api.h:6180-6187):

Returns the graph's initializers as OrtValueInfo instances. Includes constant and non-constant initializers. … For ONNX IR version >= 4, an initializer with a name that matches a graph input is considered a non-constant initializer.

A non-constant (overridable) initializer is a real, legal ONNX construct: it is the default value of an optional graph input, and the caller may replace its bytes by name at every Run. ORT even ships a tool to strip them (remove_initializer_from_input.py). So the invariant the PR asserts in its own doc comments —

  • graph_reader.rs:30 "…which lists only the constants ORT owns for the lifetime of the session"
  • ep.rs constant_input_flags doc: "ORT owns the buffer, materializes it once at session creation and cannot change it between Run calls"

— is factually wrong, and it is exactly the invariant Kernel::set_constant_inputs demands (crates/onnx-runtime-ep-api/src/kernel.rs:603 "Runtime inputs must never be marked constant: caching them would return stale results").

Consequence is a stale-result bug, not just a perf quirk. When B/scales are marked constant, MatMulNBits builds a session-lifetime packed weight once and reuses it, ignoring the current input on later calls (crates/onnx-runtime-ep-cpu/src/kernels/matmul_nbits.rs:1208-1218, self.packed_nbits_weight cache; also flips on the MLAS SQNBit path via can_prepack, matmul_nbits.rs:1092-1095 / 1728-1731). If the caller overrides B or scales at Run, the kernel returns the answer for the first B.

I reproduced it. I added a temporary test that declares B and scales as both initializers and graph inputs (an overridable initializer), created a session through the same plugin path, and read the EP's own counter. ORT printed:

[W:onnxruntime:, graph.cc:1419 Graph] Initializer B appears in graph inputs and will not be treated
  as constant value/weight. … Move it out of graph inputs if there is no need to override it …
[W:onnxruntime:, graph.cc:1419 Graph] Initializer scales appears in graph inputs and will not be
  treated as constant value/weight. …
REVIEW1110: overridable B+scales reported constant = 2 (expected 0 if sound)

ORT's runtime says these are not constant; this EP reported them constant = 2. (I reverted the temporary test; the worktree is clean.)

Reachability is not gated away: the CPU EP has no weight/graph-input claim gate (node_inputs_all_routable at ep.rs:1007 is only consulted for GPU EPs at ep.rs:392), so it claims and compiles such a node.

Fix is cheap and the correct API is already bound. ValueInfo_IsConstantInitializer exists in the vtable (OrtApi, generated in ort-sys bindings) and returns false precisely for the overridable case (onnxruntime_c_api.h:6011; doc at 6001: "an initializer that serves as the default value for a matching graph input is not a constant initializer"). read_initializer_names already iterates the OrtValueInfos — add one ValueInfo_IsConstantInitializer call per initializer and only insert names where it returns true. It reads no tensor data, so the 1 GB-weight constraint is unaffected. (ValueInfo_IsOptionalGraphInput would work equally well.)


Claims that survived falsification

Claim 2 — the mechanism (SQNBit gate + repack). Verified. can_prepack is constant_inputs[1..=4] (matmul_nbits.rs:1092) and gates mlas_sqnbit_owns_fp32_compute (:1728-1731); QLinearMatMul gates its B prepack on constant_inputs[3] (qlinear_matmul.rs:71). Leaving the flags false really does both repack per-Run and drop the faster kernel.

"Plugin EP never called set_constant_inputs". Verified. Pre-PR call sites are only onnx-runtime-session/src/executor/{bindings.rs:464,kernel_cache.rs:739}; the ep.rs:804 call is new.

Positional-vs-ORT-index alignment (the check I most expected to break). Survives. constant_input_flags is positional over node_inputs (incl. None slots); kernels index constant_inputs positionally (matmul_nbits.rs:994-996, qlinear_matmul.rs:290-292); the runtime rebuilds kernel_inputs in the same node-position order (compute.rs:2425). input_slots' per-distinct-value dedup lives only in the tensor-fetch path and never mixes with the flag array, so a node that names one value twice (Mul(x,x), shared QLinear zero-point) or has an absent optional input stays aligned. The existing a_node_that_names_one_value_twice_is_bound_once test passes.

Claim 3 — the counter test is a real falsifier, and shares statics. Verified two ways. (a) The test reads the counter from the dlopened cdylib (plugin_ort_e2e.rs:5494-5500) at the same absolute path ORT loads, so it is the same mapping/statics. (b) I reverted constant_input_flags to the producer/is_graph_input form, rebuilt, and ran it:

nbits4_decode: this EP reported 0 constant inputs, expected 2   (left: 0, right: 2)  FAILED

Exactly the message the PR predicts. This also independently proves the PR's premise that ORT hands a fused node's initializers over as subgraph graph inputs — otherwise the producer/is_graph_input form would have reported 2 and passed.

Claim 4 — second-run staleness test. Genuine, but scoped: it rotates the activation and re-runs the same session, catching an activation cached as if constant; sensitive >= 5 (:5262) keeps it non-vacuous, and ort_moved uses ORT as the oracle for detectability. Note (see minor below) it uses only pure-constant initializers, so by construction it cannot catch the weight-override bug above — a weight cache "still matches ORT."

Claim 5 — performance. Reproduced on the same host class; p50/p90, interleaved, assignment asserted before timing, inputs built once for both sides (symmetric, fair). Spot checks (ours p50):

case PR "ours after" my p50 my p90 ratio p50
bench_nbits4_m128 7.79 ms 7.64 ms 12.31 ms 7.51
bench_qlinear_u8_m1 0.079 ms 0.0885 ms 0.110 ms 1.96

The load-bearing "ours" column matches; the ours/ORT ratio drifts (7.51 vs 6.36) because ORT's p50 is the noisy side, which the PR states explicitly. Honest.

Claim 6 — read_initializer_names is cheap on a 1 GB model. Verified: it calls only GetValueInfoName, never GetTensorData; allocations are O(#initializers), not O(bytes).

Suite results against the mlas cdylib: plugin_ort_e2e 51 passed / 1 ignored; onnx-runtime-ep-plugin --lib 231 passed. Matches the PR.


Minor

  • No regression test for the override path. The three new tests never construct a non-constant initializer, and the "marks-everything-constant fails the dense=0 assertion" guard does not cover it. The temporary test I wrote (overridable B+scales ⇒ expect 0) is a ready-made regression once the flag is filtered.
  • Second-run test can't see weight staleness (plugin_ort_e2e.rs:5176-5232) — it rotates only inputs[0]. Fine for its stated goal, but it means nothing in the suite would have caught the major bug.

Nit

  • nxrt_ep_reset_constant_weight_inputs returns () but plugin_export_abi.rs:l1_required_symbols_resolve resolves it typed as fn() -> usize. Only the address is taken (never called through that type), so it's harmless and mirrors the pre-existing nxrt_ep_reset_compiled_node_count — but resolving the reset symbols as fn() would be tidier.
  • CONSTANT_WEIGHT_INPUTS is a process-global cumulative AtomicUsize (Relaxed). Correct for the lock-serialized test hook; just note it aggregates across every compile in the process and is a test-only observable with no runtime effect.

Bottom line: the diagnosis and wiring are correct and I could not break them, but the constant predicate is too broad by exactly the set ORT tells you to exclude. Gate read_initializer_names on ValueInfo_IsConstantInitializer, fix the two doc comments that assert the false invariant, and add the override regression test. Then this is a clean approve.

On the ORT plugin path every kernel saw all of its inputs as runtime
tensors, because `Kernel::set_constant_inputs` was only ever called by the
native session executor. Kernels use that flag to decide whether a prepack
may live as long as the session, and `MatMulNBits` also gates its MLAS
SQNBit path on it, so a quantized decoder rebuilt its packed weights on
every `Run` and ran a slower kernel while doing it.

The flag cannot be derived the way the routability gate derives it. ORT
hands a fused node's subgraph over with the initializers it kept inside
listed as graph inputs of that subgraph, so at Compile time a 1 GB
`MatMulNBits` `B` looks exactly like an activation. `Graph_GetInitializers`
still distinguishes them, so the graph reader now records initializer names
and the flags are keyed against that set.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@justinchuby
justinchuby force-pushed the squad/roy-plugin-constant-weights branch 2 times, most recently from 71582e2 to df733f5 Compare August 17, 2026 06:36
Review found that `Graph_GetInitializers` is documented to include
"constant and non-constant initializers". From ONNX IR version 4, an
initializer that also appears as a graph input is only a default value:
the caller may hand in a different tensor on any Run, and ORT says so at
load time ("Initializer B appears in graph inputs and will not be
treated as constant value/weight").

Keying the constant-input flags on that name list alone would let
MatMulNBits cache a session-lifetime prepack of the default weight and
return its answer forever. Filter the list through
ValueInfo_IsConstantInitializer, which reads no tensor data, and treat
an initializer as non-constant whenever ORT cannot answer: a false
negative costs a repeated prepack, a false positive costs correctness.

Adds an end-to-end regression test that builds exactly that model and
asserts both the classification (0 constant inputs) and its consequence
(both payloads match ORT run for run). Removing the gate makes it fail
with "reported 2 of them as constant weights".

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@justinchuby
justinchuby force-pushed the squad/roy-plugin-constant-weights branch from df733f5 to 0bfaad5 Compare August 17, 2026 06:37
Review follow-up: the field, the reader and the caller all say "constant
initializer" now, but the accessor was still `initializer_names()` with a
doc claiming it lists the graph's initializers. It lists only the
constant ones, and a future caller reaching for "is this value any
initializer?" would silently get the wrong answer -- the same class of
bug the previous commit fixed.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@justinchuby
justinchuby marked this pull request as ready for review August 17, 2026 07:12
@justinchuby
justinchuby merged commit 9c90674 into main Aug 17, 2026
10 of 16 checks passed
@justinchuby
justinchuby deleted the squad/roy-plugin-constant-weights branch August 17, 2026 07:47
justinchuby added a commit that referenced this pull request Aug 17, 2026
…ild shipped (#1115)

## The published wheel shipped the slow build

ORT's own CPU execution provider **is** MLAS. This repository vendors
MLAS (`crates/mlas-sys`, 833 files) and `onnx-runtime-ep-cpu` has an
opt-in `mlas` feature for it — but nothing in the packaging ever turned
it on. `python/nxrt-ep-cpu/setup.py` ran `cargo build --release -p
onnx-runtime-ep-cpu-plugin` with `CARGO_FEATURES: list[str] = []`, so
**every published `nxrt-ep-cpu` wheel contained the pure-Rust fallback
paths**.

That is not a small difference. Measured end-to-end through the plugin
path (the harness added in #1110), same host, same ORT, interleaved A/B:

| case | ours/ORT p50, no MLAS | ours/ORT p50, MLAS | ours p50 (ms) no
MLAS → MLAS |
|---|---|---|---|
| `MatMulNBits` int4 M=128 | **81.1x** | 7.3x | 115.9 → 8.80 |
| `MatMulNBits` int4 M=1 | 14.8x | 5.2x | 1.850 → 0.400 |
| `MatMulNBits` int4 f16-act M=1 | 20.1x | 4.7x | 1.418 → 0.437 |
| `MatMulNBits` int8 M=256 | 7.1x | 1.5x | 46.38 → 13.86 |
| `QLinearMatMul` u8 M=128 | **54.9x** | 9.3x | 47.62 → 10.36 |
| `QLinearMatMul` u8 M=1 | 122.8x | 2.1x | 6.414 → 0.092 |
| `QLinearMatMul` i8 M=1 | 10.1x | **0.038x** | 5.457 → 0.092 |
| `MatMul` f32 M=128 | 1.58x | **0.82x** | 12.56 → 3.82 |
| `MatMul` f32 M=1 | 30.9x | 1.32x | 5.246 → 0.118 |
| `MatMul` f16 M=1 | 0.58x | 0.85x | 2.321 → 2.084 |
| `MatMul` f16 M=128 | 2.73x | 4.08x | 24.76 → 20.36 |

Host: AMD EPYC 9V74 (32 vCPU / 16 physical cores, AVX2+FMA+F16C, no
AVX-512), ORT 1.27.0, release build, K=N=2048, 3 warmups + 41
interleaved iterations, p50. Ratios are **ours/ORT**, so below 1.0 means
we are faster. Cold session-creation time is reported separately by the
same harness and is not folded into these numbers. The two f16 rows move
in opposite directions because that path does not go through MLAS at all
— they are the control, and they bracket the host noise for this table
(±0.3x at p50 on a shared machine).

## What this PR changes

1. **`python/nxrt-ep-cpu/setup.py` enables the feature** on every
`(system, machine)` pair in `MLAS_TARGETS`. `NXRT_EP_CPU_NO_MLAS=1`
builds the pure-Rust cdylib anyway, for a toolchain with no C++
compiler.

2. **`onnx-runtime-ep-cpu-plugin` gets its own `mlas` feature** that
forwards to `onnx-runtime-ep-cpu/mlas`, so the cdylib crate can `cfg` on
it.

3. **The cdylib now exports `nxrt_ep_build_features()`.** A compiled
library says nothing about how it was built, and here that difference is
81x. `setup.py` refuses to bundle a cdylib whose report disagrees with
what it asked cargo for (`target/release` is shared with every other
build in the checkout, so the file that exists after `cargo build` is
not necessarily the file that build produced), and the wheel's
cibuildwheel smoke test re-checks the installed artifact against
`nxrt_ep_cpu._build.EXPECTED_FEATURES`, which `setup.py` generates.
`nxrt_ep_cpu.build_features()` exposes the same fact to users.

4. **A test-harness bug that made feature-specific testing
meaningless.** `onnx_runtime_ort_testkit::find_plugin_cdylib` rebuilds
the cdylib with `cargo build -p <pkg>` — *no features*. So `cargo test
-p onnx-runtime-ep-cpu-plugin --features mlas` compiled the test binary
with MLAS and then **overwrote the MLAS cdylib with a default-feature
build**, asserting against the wrong library, silently, in the direction
that hides problems. `find_plugin_cdylib_with_features` fixes it and
`cdylib_resolve.rs` passes the features it was compiled with.
Consequence worth stating: the ORT conformance suite has now run against
the MLAS cdylib for the first time, and passes on both feature sets.

5. **CI builds the MLAS cdylib on each lane that matches a wheel
target** — `Fast (Linux x86_64)`, `Rust coverage (Windows x86_64)`,
`Rust coverage (macOS arm64)`, `Rust (Windows ARM64)`. A target is only
listed in `MLAS_TARGETS` if a lane compiles it; if these lanes go red on
some platform I will remove that platform from the set rather than ship
a wheel that fails to build at release time.

## Tests

- `l1_build_features_match_the_compiled_feature_set` (new,
`plugin_export_abi`) — dlopens the cdylib the harness resolved and
asserts its reported features equal `cfg!(feature = "mlas")`. **This is
the falsifier for item 4:** with the testkit fix reverted it fails with
`the cdylib at …/libonnx_runtime_ep_cpu_plugin.so reports features ""
but this test binary was built with "mlas"`.
- `check_wheel.py` (new) replaces the wheel's one-line `test-command`.
Falsified by hand: editing the installed `_build.py` to claim
`"avx9000"` fails with `bundled cdylib reports build features 'mlas',
but this wheel was built asking for 'avx9000'`.
- `plugin_export_abi::l1_no_symbol_leakage` — the new export is added to
the allow-list; `nm -D` on the MLAS cdylib shows the same 8 exported
symbols as before plus this one, i.e. the vendored C++ does not leak
symbols.
- Local wheel build + install + smoke test on Linux x86_64: `OK
…/libonnx_runtime_ep_cpu_plugin.so features='mlas'`.

## Verification

- `cargo test -p onnx-runtime-ep-cpu-plugin` — 50 + 9 + 7 + 1 passed
(default features)
- `cargo test -p onnx-runtime-ep-cpu-plugin --features mlas` — same
counts, all passed (first time this actually tested the MLAS build)
- `cargo clippy --all-targets` clean for `onnx-runtime-ep-cpu-plugin`
and `onnx-runtime-ort-testkit`, with and without `--features mlas`
- `cargo fmt --all -- --check` clean
- `python -m build --wheel` + install + `check_wheel.py` green locally

## Still losing after this change

With MLAS, on this host: `MatMul` f16 M=128 4.08x, `MatMulNBits` int4
M=128 7.3x and M=1 5.2x, `QLinearMatMul` u8 M=128 9.3x. Those are kernel
gaps and stay open on my task list — this PR only stops us from shipping
the *much* slower build. Nothing here reduces precision or hides setup
cost.


---

## Post-review fix (commit 2): the smoke test pointed at the wrong path

Independent review found a release-blocking bug, and it was right.

`test-command = 'python {project}/check_wheel.py'` — but cibuildwheel is
invoked as `cibuildwheel python/nxrt-ep-cpu` **from the repository
root** (`publish-ep-plugins.yml:122`), so `{project}` is the repository
root and `{package}` is this directory. All four wheel lanes would have
failed with `python: can't open file '/project/check_wheel.py'`, and
only on an `nxrt-ep-v*` tag — the release-time failure this PR exists to
prevent.

Fixed to `{package}`, and pinned by a new test binary that runs in
ordinary CI:

- `wheel_test_command_names_a_file_that_exists` — rejects `{project}`
and asserts the referenced script exists. Falsified by restoring
`{project}`: *"test-command uses {project} (the repository root)"*.
- `every_mlas_wheel_target_is_built_by_a_ci_lane` — asserts each
operating system in `MLAS_TARGETS` has a lane in `ci.yml` that compiles
the MLAS cdylib. Falsified by adding `("freebsd", "x86_64")`: *"enables
MLAS for the operating systems {"darwin", "freebsd", "linux", "windows"}
but ci.yml builds the MLAS cdylib on only 3 lanes"*. It counts operating
systems rather than targets because one coverage-matrix step covers both
`windows/amd64` and `darwin/arm64`.

Also from the review: the `comma-separated` wording in the
`nxrt_ep_build_features` doc (only one token is ever emitted), the
undocumented `AttributeError` in `build_features()`, the stale
`_mlas_features()` reference in the pyproject comment, and a note on the
testkit cache key (feature sets share one `target/<profile>` path; no
caller resolves two in one process, and the key stops the two answers
being conflated).

Rebased onto `main` after #1110; the export allow-list now carries both
that PR's counters and this PR's build-identity symbol.

### Verification (re-run after rebase)

- `cargo test -p onnx-runtime-ep-cpu-plugin --features mlas` — 52 e2e (1
ignored) + 9 + 7 + 2 + 1 passed
- `cargo test -p onnx-runtime-ep-cpu-plugin` — same counts on default
features
- `cargo fmt --all -- --check` clean **for this branch**; note `main` is
currently fmt-broken at
`crates/onnx-runtime-ep-cuda/src/kernels/device_argmax.rs`, repaired by
#1118
- `cargo clippy --all-targets` clean for both crates, with and without
`--features mlas`


---

## Second review round (commit 3): the CI guard could pass while broken

Review returned APPROVE with two MINORs that were both real, and both
are fixed.

**1. `every_mlas_wheel_target_is_built_by_a_ci_lane` was vacuous under
the exact failure it guards.** It counted text occurrences of the MLAS
build command in `ci.yml` and compared against the number of operating
systems. darwin/arm64's only MLAS build is the `if: runner.os !=
'Linux'` step on the coverage matrix, so **deleting the `macos-latest`
matrix row removes that build while the count stays at 3** and the test
stays green.

It is now structural: it parses `ci.yml`, resolves each job's runner
operating systems (expanding `runs-on: ${{ matrix.os }}` over
`matrix.include[].os` and plain `matrix.os`), applies each step's
`runner.os` condition, and asserts the union covers every operating
system in `MLAS_TARGETS`. Both falsifiers fire:

- delete the `macos-latest` row → *"setup.py enables MLAS for {"darwin",
"linux", "windows"} but no ci.yml lane compiles the MLAS cdylib on
["darwin"] … (lanes cover {"linux", "windows"})"*
- add `("freebsd", "x86_64")` → the same, for freebsd.

**2. CI built the MLAS cdylib but never tested it.** With default
features the build-identity assertion is trivially satisfied (`"" ==
""`), the vendored C++ is not linked so the leakage check has nothing to
leak, and the testkit rebuild defect this PR fixes is only observable
when features are requested — so reverting that fix would have left CI
green. The ORT-gate lane (the only one with a real ONNX Runtime) now
runs the plugin suite a **second time** with `--features mlas`, which is
where the conformance and build-identity claims are actually enforced.

**NIT:** both ctypes identity reads now release the handle in a
`finally` (on Windows a retained handle locks the DLL for the life of
the process).

**NIT — thread counts, which the tables omitted.** Every measurement
below and above uses **default `SessionOptions`** on both sides: the
harness never sets `intra_op_num_threads`, so ORT uses its own default
(physical cores — 16 on this host) and our EP uses its own pools (sized
from available parallelism — 32 vCPUs). Both sides therefore run
multi-threaded, and neither is throttled. Ratios are ours/ORT p50, so >1
means we are slower.

### Verification (re-run after rebase onto `main` @ b7fa5e1)

- `NXRT_REQUIRE_ORT_TESTS=1 cargo test -p onnx-runtime-ep-cpu-plugin
--features mlas` — 52 e2e (1 ignored) + 9 + 7 + 2 + 1 passed
- `NXRT_REQUIRE_ORT_TESTS=1 cargo test -p onnx-runtime-ep-cpu-plugin` —
identical counts on default features
- `cargo clippy --all-targets` clean for both crates, with and without
`--features mlas`
- `cargo fmt --all -- --check` clean (`main`'s unrelated fmt breakage
was repaired by #1120; my #1118 was closed as superseded)

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant