Repository navigation
fix(ep): claim and correctly execute MatMulNBits and QLinearMatMul - #1101
Merged
Merged
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1101 +/- ##
==========================================
+ Coverage 80.17% 80.56% +0.39%
==========================================
Files 367 367
Lines 155918 156610 +692
Branches 155918 156610 +692
==========================================
+ Hits 125004 126170 +1166
+ Misses 26232 25730 -502
- Partials 4682 4710 +28
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
|
justinchuby
force-pushed
the
squad/roy-matmul-claim-bugs
branch
from
August 17, 2026 02:57
dc37bd1 to
ea13281
Compare
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
## What ``` $ git checkout origin/main && cargo fmt --all -- --check Diff in crates/onnx-runtime-ep-cpu/src/kernels/governed_weight_cache.rs:263 Diff in crates/onnx-runtime-ep-cpu/src/kernels/matmul.rs:354, 2127, 2148, 2163, 3206 Diff in crates/onnx-runtime-ep-cpu/src/lib.rs:92 ``` `main` is fmt-broken at a0ed99a (#1100). `main` is unprotected and the merge queue does not re-run CI on the merge result, so this fails the *Check formatting* step of both **Fast (Linux x86_64)** and **Rust quality** on every branch cut from it — including #1101, which is how I found it. ## What this is The literal output of `cargo fmt --all`, nothing else. No logic change, no test change. 17 insertions, 20 deletions, all whitespace and line-wrapping. ## Verification - `cargo fmt --all -- --check` — clean - `cargo test -p onnx-runtime-ep-cpu --lib` — unchanged Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Four independent defects meant the plugin EP had never executed either
quantized matmul under an ORT host, whatever the assignment policy said.
1. `supported_dtypes_for_op("MatMulNBits", "com.microsoft")` returned
`F32_ONLY`, but the node filter requires *every* input and output dtype
to be listed and the packed weight is `uint8`. Every int4/int8 node was
dropped and ORT's own kernel ran.
2. `ShapeInference::for_node` aliased `MatMulNBits` to plain `MatMul`, which
broadcasts the activation against the *packed* `[N, blocks, blob]` weight:
K=256/N=4096 inferred `[4096, 1, 16]` instead of `[1, 4096]`.
3. `QLinearMatMul` was absent from both shape tables, so it resolved to
`Declined` and the fail-closed shape filter dropped the claim. Its
operands are inputs 0 and 3, not 0 and 1.
4. `input_slots` numbered ORT inputs by position. ORT's fused-node metadata
carries a *set* of input names, so a node naming one value twice — a
`QLinearMatMul` sharing a zero-point or scale initializer, or any
`Mul(x, x)` — shifted every later slot off the end of the bound array.
The single-kernel fast path now fails closed with a diagnosable message
instead of indexing out of bounds and surfacing as `Compute: internal panic`.
Falsified by a 10-case generated-model conformance sweep that asserts each
op lands on our EP with CPU fallback disabled and then matches a plain-ORT
session elementwise, plus hardware-independent shape and dtype unit tests.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Review follow-up. `supported_dtypes` is the *union* of an op's edge dtypes, so on a block-quantized op it admits combinations no kernel accepts: a `MatMulNBits` with `float16` `zero_points` — spec-valid, and a session ORT builds happily — passed the union test because both `float16` and `uint8` are somewhere in the union, and then failed inside `execute`. Adds `KernelRegistryEntry::input_dtype_constraints`, a per-input-slot override that the node filter applies instead of the union, populated for `MatMulNBits` (float compute on A/scales/bias, uint8 on B/zero_points, int32 on g_idx) and `QLinearMatMul` (float32 scales, uint8/int8 operands and zero points) directly from those kernels' `require_dtype` calls. Every other op keeps the union, which is exact for a uniform op. Also: - `Mul(x, x)` falsifier for the input-slot dedup. ORT dedups *graph inputs* too, not only initializers, so the old positional mapping ran off the end for any node naming one value twice: `input slot 1 maps to ORT input 1, but ORT bound only 1 input(s)`. - A regression test pinning that a float16-zero-point `MatMulNBits` is declined and still computed, rather than claimed and failed. - Unit test for the per-slot filter; corrected a line-wrapped error message and two stale doc comments. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
squad/roy-matmul-claim-bugs
branch
from
August 17, 2026 03:14
ea13281 to
905f7c4
Compare
`ONNXTensorElementDataType` is `c_uint`, which bindgen renders as `u32` on Linux and `i32` on Windows, so hardcoding `u32` broke the Windows x86_64 and ARM64 lanes at compile time. Use the alias. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
marked this pull request as ready for review
August 17, 2026 04:37
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
`main` (#1101) independently hit the same dtype-filter bug class for `MatMulNBits`/`QLinearMatMul` and added `FLOAT_COMPUTE_DTYPES` — the identical f32/f16/bf16 set as this branch's `MLAS_FLOAT_DTYPES`. Resolved `kernels/mod.rs` by taking main's block verbatim, dropping the duplicate constant, pointing `Conv` at `FLOAT_COMPUTE_DTYPES` and folding this branch's rationale into main's doc comment. The merge also surfaced two shape-table drifts, both caught by `every_registered_op_has_a_shape_rule_or_is_a_known_gap`: * The probe supplied no attributes, so `com.microsoft::MatMulNBits` — whose rule derives the output width from `N` — looked like a gap it is not. A node without `N` is malformed rather than defaulted, so no real graph hits that path. The probe now sweeps an attribute bundle alongside the empty one, keeping default-fallback rules honest while modelling production for rules that cannot default. * With that fixed, the first assert stopped masking the second: `QLinearMatMul` gained a rule in main and was stale in `DECLINED`. Removed. Gap count 65 -> 64, group 3 35 -> 34. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
`main` is currently red on `cargo fmt --check`: #1101 landed one slice expression that stable rustfmt reformats. That fails the "Check formatting" step in both `Fast (Linux x86_64)` and `Rust quality`, on `main` and on every open PR. This applies **exactly the diff CI printed**, verbatim: ```diff for c in 0..group { - let block_values = &packed_rows[c] - [block * packed_block_size..(block + 1) * packed_block_size]; + let block_values = + &packed_rows[c][block * packed_block_size..(block + 1) * packed_block_size]; ``` No behaviour change, no other file touched. `cargo fmt --all -- --check` is clean locally after it. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
added a commit
that referenced
this pull request
Aug 17, 2026
## What this changes When our CPU EP is selected, it must not hand work to ORT's CPU EP. This PR removes the **three** mechanisms by which it was doing so for the activation and normalization families. `GetCapability` runs three independent fail-closed filters and a claim must clear all of them; each of the last two was found only after this PR had already claimed the job was done. ### 1. The performance-based decline policy is deleted `assignment_policy.rs` (2118 lines) measured whether we beat ORT on a given shape/dtype and, where we lost, returned `ClaimPreference::defer` so ORT's CPU EP would take the node. That whole file is gone, along with the `claim_preference` override in `provider.rs`. `claim_preference_node` now returns `Claim` immediately. Beyond the architectural rule, the policy could not have worked as intended: - **A deferral splits the graph.** Every declined node is a partition boundary, which costs fusion, prepacking and buffer reuse across it — none of which the per-node threshold accounted for. - **It cannot see the thread count.** Capability runs before the session's intra-op pool is known. `Sqrt` at 64 Ki wins 1.9x against a single-threaded ORT and loses at 0.30x against 16 threads. One number cannot be right for both. - **Sometimes there is nothing to defer to.** ORT has no bf16 kernel for these ops and no f16 kernel for most. Declining a bf16 `Gelu` does not get a faster kernel, it gets a load failure. - **The thresholds were tuned on one host.** Every number was measured on a single EPYC 9V74. Shipping it made every other machine's latency a guess. Removing the override is also a small capability-time win: the default adapter deep-clones every input `Shape` and collects dtypes for every node in order to build the metadata the policy consumed. ### 2. The shape-inference filter was declining the same ops anyway Deleting the policy did not, by itself, achieve the goal. `GetCapability` runs a second, independent fail-closed filter (`onnx-runtime-ep-plugin/src/ep.rs`): it drops any claim containing a node whose `ShapeInference::for_node` returns `Declined`, and that match ends in `_ => Declined`. **An op we register a kernel for, but which is absent from that table, is silently handed to ORT no matter what `supports_op` answers.** This is the same mechanism that made the `com.microsoft` activations unreachable until #1082. The trigonometric, hyperbolic and remaining activation ops were all in that gap. Now listed: `Sin`, `Cos`, `Tan`, `Asin`, `Acos`, `Atan`, `Sinh`, `Cosh`, `Asinh`, `Acosh`, `Atanh`, `ThresholdedRelu`, `Swish`, `com.microsoft::Silu` and `PRelu`. `Silu` had been deliberately excluded with the comment that ORT has no kernel for it. That reasoning was backwards: an op ORT cannot run is precisely the one we must never hand over. `GroupNormalization` was in the same gap and is now covered too. ### 3. The dtype filter was declining a different set of ops Review found a third filter, and it was still handing over one of the very ops section 2 had just fixed. `node_passes_dtype_filter` looks the node's op up in the plugin's `KernelRegistryEntry` list and returns `false` when there is no entry. That list is built from `build_cpu_registry_with_descriptors`, which recorded keys as they were registered — but `register_cnn_ops` takes `&mut OpRegistry` and writes *past* the recording wrapper. Eighteen ops were in the registry and absent from the descriptors, so `supports_op` claimed each one and capability then dropped it: > `PRelu`, `BatchNormalization`, `InstanceNormalization`, `GroupNormalization`, > `Conv`, `ConvTranspose`, `MaxPool`, `AveragePool`, `GlobalMaxPool`, > `GlobalAveragePool`, `GlobalLpPool`, `LpPool`, `Resize`, `GridSample`, > `AffineGrid`, `Col2Im`, `CenterCropPad`, `SpaceToDepth` Four are activations or normalizations this EP owns. **`PRelu` is the sharpest case: section 2 gave it a shape rule, so it cleared filter two, and the dtype filter declined it anyway.** Both pure-Rust inventory tests passed while real ORT ran the node. Descriptors are now derived from `OpRegistry::keys()` instead of a parallel recorded list, making the two sets identical by construction rather than by convention. They are also sorted: they get leaked into a `'static` slice ORT reads, and hash-map iteration order would make any snapshot diff flap. `descriptors_derived_from_real_registry_not_hand_maintained` had been asserting this bug as *correct behaviour* — it allowed a delta of up to 50 entries and named CNN ops as the expected difference. It now asserts set equality. The lesson, having now been caught twice: an inventory test is only as good as its source of truth, and two review rounds passed on tests that enumerated the wrong set. The only check that cannot be fooled this way is the end-to-end one that asks real ORT which EP got the node. ## Scope — what this does *not* fix **64 registered ops remain in the shape-inference gap.** This PR closes the activation/elementwise families; it does not close the gap universally. The full list is asserted exactly by the new inventory test and summarised in `docs/performance/CPU_ACTIVATION_GAPS.md`: - **20 data-dependent** — output shape is a function of an input's *values* (`NonZero`, `Unique`, `Compress`, `Expand`, `Tile`, `Pad`, `TopK`, `Split`, `Unsqueeze`, `Resize`, `AffineGrid`, `Col2Im`, `CenterCropPad`, ...). Correctly declined today, though most carry a constant initializer in practice, so a pass that resolves initializer values at capability time could claim them. - **10 internal fusion ops** created after capability, never candidates (`FusedGemm`, `FusedAttention`, `FusedMatMulBias`, the `pkg.nxrt` ops). - **34 inferrable but unwritten — this is the work.** Ten are one-line shape-preserving rules (`QuantizeLinear`, `DequantizeLinear`, `CastLike`, `ScatterND`, `Trilu`, `CumSum`, ...); nine are pooling/CNN geometry (`MaxPool`, `AveragePool`, `Global*Pool`, `ConvTranspose`, `GridSample`, `SpaceToDepth`) inferrable exactly as `build_conv` already does for `Conv`; eight more follow from attributes (`ArgMax`, `Flatten`, `GatherElements`, `Size`, ...); the rest are contrib and model ops. Two entries deserve singling out. `com.microsoft::Attention` is *the* attention op in exported GenAI models, and the existing opset-23 arm is guarded to the default domain, so we hand it over. `LinearAttention` (both domains) and `com.microsoft::CausalConvWithState` are the Qwen3.5 / Qwen3-Next hybrid linear-attention primitives — **ORT has no kernel for them at all**, so declining them does not get a faster implementation, it gets a load failure. Closing that remainder is the next PR. It is a different domain from the activation kernels and each entry needs its own numeric test. ### Corrections to my earlier figures I published two wrong counts before this test existed, in opposite directions, and the test found both. **66, then 52, now 66 again — for different reasons each time.** The first figure came from a scratch script that text-matched the table and missed its guard arms, so it reported the whole `Reduce*` family (handled by `op_name if is_reduction(op_name)`) as a gap. Correcting that, I over-corrected and claimed the pooling family "has rules already" — it does not; `compute.rs` has no pooling arm at all. The 52 figure was also built on `build_cpu_registry_with_descriptors`, which is **not** the registry: `register_cnn_ops` writes straight to the inner `OpRegistry`, so 14 CNN ops and `PRelu` never appear in the descriptors at all. The test now enumerates `OpRegistry::keys()` — the same set `supports_op` consults. **`MoE`/`QMoE`/`LinearAttention`/`CausalConvWithState` were misclassified as internal fusion ops.** They are read from exported models — `deepseek_v2_tiny_qmoe_native_e2e.rs` asserts a loaded graph contains a `com.microsoft::QMoE` node, and the linear-attention pair are Qwen3.5 primitives. All are real gaps. **Three gaps I had missed entirely:** `Unsqueeze` (declines whenever `axes` is input[1], i.e. every opset-13+ graph), `com.microsoft::Attention`, and `EyeLike`. This is the argument for the inventory test: hand-maintained prose about which ops reach ORT was wrong three times in a row, in both directions, and each time it read as confident. ## Tests ### Inventory - the test that would have caught this class of bug `every_registered_op_has_a_shape_rule_or_is_a_known_gap` enumerates `OpRegistry::keys()` and asserts the set of ops that decline shape inference *exactly*. Registering an op without a shape rule fails; adding a shape rule without removing the op from the list also fails. Neither direction can pass silently, and the allowlist doubles as the gap inventory. It needs no ORT, so it runs in every job rather than only the ORT-gated one. The probe is a sweep, not a point: opsets {1, 13, 18, 22, 23} x arities 1..4 x ranks 1..4, counting an op as declining only when nothing in the matrix produces a rule. A single one-input rank-2 probe reported `Conv` as a gap, because `build_conv` reads `input_shapes[1][0]` for its output channel count and needs rank >= 3. Sweeping removes that artifact and keeps the allowlist from encoding one arbitrary opset. Verified both directions falsify: deleting `| "Cos"` fails with `[("", "Cos")]`; adding `| "Trilu"` fails with `these ops now have a shape rule but are still listed as declined: [("", "Trilu")]`. `no_activation_or_norm_op_is_left_to_ort` is a standing guard on the 39 activation and norm ops this EP owns - the families #1082, #1093 and #1097 made reachable. It asserts in both directions: a filter of the form `registered.contains(op) && declines(op)` would let a *deleted registration* pass silently, which is the same hand-off by a different route. Verified by renaming the `Silu` kernel key, which now fails with `these ops are no longer registered by the CPU EP, so ORT will execute them: [("com.microsoft", "Silu")]`. That direction immediately caught two entries I had wrong: `Swish` is registered in the default domain, not `com.microsoft`, and `HardSwish` has no kernel at all. ### Assignment sweeps Two sweeps replace the 13 deleted deferral tests. Both iterate 21 fixtures: - `no_supported_node_is_ever_left_to_the_ort_cpu_ep` - for every fixture, the op under test appears in *our* EP's node list and in no other EP's. - `every_fixture_loads_with_cpu_fallback_disabled` - loads each fixture with `session.disable_cpu_ep_fallback`, so a silent hand-off becomes a load failure rather than a slow success. Verified the sweep falsifies too: with `| "Sin"` removed from the table it fails with `[sin_assignment_f32] ours=[], others=["Sin"]`. Both run fail-closed in CI: `conformance_setup` panics rather than skipping when `NXRT_REQUIRE_ORT_TESTS=1`, which the `CLI ORT` job sets. (That job's plugin step was itself being skipped whenever an earlier step failed - fixed separately in #1096.) `erf_reference` would have become dead code when the deferral tests were deleted. Rather than remove it, it is now used by `float16_biasgelu_runs_on_our_ep_with_correct_numerics`, which checks *our* kernel's numerics - more important now that we always execute `BiasGelu` instead of sometimes declining it. ## Docs `CPU_MATMUL_ASSIGNMENT.md` is reframed from claim/defer to win/gap. Every measurement is unchanged — the numbers still say exactly where we are slower than ORT; they are now a work list rather than a decline table. `CPU_ACTIVATION_GAPS.md` is new: every range where we still lose, and the two root causes. Neither is polynomial accuracy: 1. **Our elementwise kernels are single-threaded** while ORT splits across its intra-op pool. This is the flat ~0.7-0.8x plateau across the f32 activation family, and it is worth more than any further approximation work. `KernelContext_ParallelFor` is the untried lever. 2. **~1.2 us of fixed per-node plugin dispatch overhead**, which is what the ~0.75-0.8x at n=1 measures. ## Validation `cargo fmt`, clippy clean on the three affected crates, 1266 `ep-cpu` lib tests, 224 `ep-plugin` unit tests, 2 inventory tests, and 37 plugin E2E tests under `NXRT_REQUIRE_ORT_TESTS=1` with real ORT. ## Update — third decline path (commit `67a611353`) Opus review round 4 returned **NO-GO** on the grounds that a fourth decline path existed and the PR's central claim was therefore still false. It was right. See section 3 above. Independently verified before fixing: 18 of the 177 unique `(domain, op_type)` pairs in the registry had no descriptor, including `PRelu`, `BatchNormalization`, `InstanceNormalization` and `GroupNormalization`. Counting individual registry entries (`op_type` + `domain` + `since_version`, the unit `OpRegistry::len()` reports) the registry holds 208, and descriptors now match it exactly — `descriptors_derived_from_real_registry_not_hand_maintained` asserts that equality. New guards, each verified to fail when its invariant is broken: | guard | falsified by | observed failure | | --- | --- | --- | | `every_registered_op_has_a_kernel_registry_entry` | filtering `PRelu` out of descriptors | `1 registered ops have no kernel-registry entry ... ["::PRelu"]` | | `activation_and_norm_ops_clear_every_capability_filter` | same | `::PRelu: no kernel-registry entry (dtype filter declines it)` | | `prelu_assignment_f32` (real ORT) | same | `ours=[], others=["PRelu"]` → `'PRelu' must run on this EP` | | `descriptors_derived_from_real_registry_not_hand_maintained` | same | names the missing ops instead of tolerating a delta of 50 | After the fix, real ORT reports `ours=["PRelu"], others=[]` and `ours=["GroupNormalization"], others=[]`. Writing `activation_and_norm_ops_clear_every_capability_filter` immediately found two more real gaps: **`Celu` and `Mish` have no kernel at all**. That is a missing feature rather than a decline, so they are excluded from that test with a comment naming them, and recorded in `CPU_ACTIVATION_GAPS.md` under a new "Activations with no kernel at all" section rather than quietly dropped. Re-validated: `cargo fmt --all --check`, scoped clippy clean, 1267 `ep-cpu` lib tests, 224 `ep-plugin` unit tests, and 57 plugin tests under `NXRT_REQUIRE_ORT_TESTS=1` (37 E2E + 4 coverage + 9 + 6 + 1). ## Update — Opus round 5: `GO WITH FINDINGS` Round 5 confirmed the central claim now holds, verified against real ORT (`ours=["PRelu"], others=[]` and `ours=["GroupNormalization"], others=[]`), and traced every node-removing gate in `ep_get_capability_inner` to confirm no fifth decline path exists for the activation/norm families. Two findings, both addressed: **Finding 1 (minor, real) — `Conv` advertised a dtype its kernel rejects.** Giving `Conv` a `KernelRegistryEntry` was itself an over-claim: `supported_dtypes_for_op("Conv", "")` returned `FLOAT_DTYPES`, which includes f64, but `ConvKernel::execute` rejects anything outside f32/f16/bf16. Before this PR `Conv` had no descriptor so an f64 `Conv` was declined; after it, the node would clear both filters, compile, and then fail at `Run`. Fixed by adding `MLAS_FLOAT_DTYPES` and giving `Conv` its own arm. This is not a re-introduced fallback. Advertising a dtype we cannot execute is a different thing from declining one we can: f64 `Conv` is genuinely unsupported, and reporting that honestly at capability time is correct. The rest of the CNN family really does dispatch f64 through `dispatch_float!` and keeps `FLOAT_DTYPES`. Pinned by `conv_does_not_advertise_a_dtype_its_kernel_rejects`. **Finding 2 (minor, docs) — stale counts in this description.** The breakdown said 20 + 10 + 36 = 66 while the total said 65, and still listed `GroupNormalization` among the unwritten shape-preserving rules even though this PR wrote one. Corrected to 35 above. The "177" figure was unique `(domain, op_type)` pairs, not `OpRegistry::len()` (208) — both are now stated explicitly. The committed `CPU_ACTIVATION_GAPS.md` was already correct. --- ## Merge with `main` `main` moved under this branch and #1101 independently hit the *same* dtype-filter bug class, for `MatMulNBits` and `QLinearMatMul`, adding `FLOAT_COMPUTE_DTYPES` — byte-for-byte the same f32/f16/bf16 set as this branch's `MLAS_FLOAT_DTYPES`. Two people finding the same trap independently is the strongest evidence that the fail-closed dtype filter needed the systematic fix in this PR rather than another per-op patch. `kernels/mod.rs` was resolved by taking main's block verbatim, deleting the duplicate constant, pointing `Conv` at `FLOAT_COMPUTE_DTYPES` and folding this branch's rationale into main's doc comment. The merge then made the inventory test earn its keep twice: **The probe was lying about `com.microsoft::MatMulNBits`.** `declines()` supplied no attributes, deliberately: an attribute-dependent rule should fall back to the ONNX default when the attribute is absent, and the empty bundle keeps that honest. But `MatMulNBits` derives its output width from `N`, and a node without `N` is *malformed*, not defaulted — `for_node` declines it, correctly. No real graph reaches that path, so the test was reporting a gap that does not exist. The probe now sweeps a plausible attribute bundle **alongside** the empty one, so default-fallback rules are still checked against absence while rules that cannot default are modelled the way production sees them. **With that fixed, the first assert stopped masking the second.** `QLinearMatMul` gained a real rule in main and was stale in this branch's `DECLINED` list. Removed. That is the drift check doing exactly the job it was written for — the list cannot quietly stop describing reality in either direction. Gap count 65 → 64; group 3, 35 → 34. `CPU_ACTIVATION_GAPS.md` updated to match. Post-merge validation, all green: `cargo fmt --all -- --check`; scoped clippy (`onnx-runtime-ep-cpu`, `-ep-plugin`, `-ep-cpu-plugin`, `--all-targets`) clean; 1280 `ep-cpu` lib tests; and the full plugin suite under `NXRT_REQUIRE_ORT_TESTS=1` — 40 real-ORT E2E assignment tests plus the 4 inventory tests. > The `Fast` lane on this branch will stay red on `cargo fmt --all -- --check` > until #1107 lands. That failure is `main`'s, from #1101, in > `matmul_nbits.rs` — a file this PR does not touch. The fix is deliberately > **not** duplicated here; this branch will pick it up by merging `main` once > #1107 is in. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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.
The finding
Our CPU EP has never executed
MatMulNBitsorQLinearMatMulthrough the ORT plugin path. Not "deferred them under policy" — never claimed them at all. Every int4/int8 quantized model that has run through this EP had its matmuls executed by ORT's own kernels, silently, regardless of whatassignment_policysaid and regardless ofsession.disable_cpu_ep_fallback.That invalidates the reachability half of every plugin-path claim in
CPU_MATMUL_ASSIGNMENT.md. It also means the recent architectural direction — this EP must never hand a node to ORT's CPU EP — could not have been satisfied by removing the defer matrix alone.There are four independent defects, each of which is sufficient on its own to lose the op.
1. The dtype filter excluded
MatMulNBitsoutrightnode_passes_dtype_filter(ep.rs) returnsfalseunless every input and output dtype of the node appears inentry.supported_dtypes.supported_dtypes_for_op("MatMulNBits", "com.microsoft")returnedF32_ONLY.MatMulNBitsinputs areA: float,B: uint8,scales: float,zero_points: uint8,g_idx: int32. Theuint8weight fails the filter on input 1, so the claim is dropped before anything else is consulted. The correct set is the one the kernel itself enforces atmatmul_nbits.rs:780-815: float32/float16/bfloat16 onA/scales/Y,uint8onBandzero_points,int32ong_idx.2.
MatMulNBitsused the plain-matmul shape ruleShapeInference::for_nodemapped"MatMul" | "MatMulNBits"toSelf::MatMul.Bis not a matmul operand — it is the packed[N, blocks_per_col, blob_bytes]weight — so the rule broadcast the activation against the packing:[4096, 1, 16][1, 4096][8, 1, 16][1, 8]The correct rule is the activation shape with its last dim replaced by the
Nattribute, which is whatonnx-runtime-shape-inference'squantized_matmul(handlers/linalg.rs:36) already does. BecauseNis only visible on the node, the attribute-freefor_op_domaintable must decline rather than fall back to the matmul rule.Defect 1 masked defect 2: fixing only the dtypes turns a silent escape into a hard
Y must have shape [1, 4096], got [4096, 1, 16].3.
QLinearMatMulwas in neither shape tableMissing from
for_op_domainandfor_node, so it resolved toShapeInference::DeclinedandGetCapability's fail-closed shape filter dropped the whole claim. This is the same failure mode as thecom.microsoftactivations fixed in #1082.Note that the
MatMulrule would have been wrong here too: the ONNX operand order isa, a_scale, a_zero_point, b, b_scale, b_zero_point, y_scale, y_zero_point, so the operands are inputs 0 and 3.4.
input_slotsnumbered ORT inputs by position, not by valueOnce claimed,
qlinear_u8still failed:ORT's fused-node metadata carries a set of input names. A node that names the same value twice is bound once, and the sequential numbering shifted every later slot — the last one past the end of the array ORT actually passes. This is not exotic: a real quantized
QLinearMatMulroutinely shares one zero-point or scale initializer across itsa/b/ytriples, and anyMul(x, x)-shaped node does the same. Onmainthis surfaced as an opaqueCompute: internal panicfrom an out-of-bounds index across the C ABI.Slots are now assigned per distinct value in first-appearance order, and the fast path fails closed with the message above instead of panicking.
5. The union dtype list over-claimed what the kernels accept
Found by review of the first commit.
supported_dtypesis a flat set and the node filter tests membership in it, so on a mixed-dtype op it is strictly weaker than the kernel's own rule.com.microsoft::MatMulNBitsspec-allowszero_pointsinuint8,int32,float16orfloat(contrib_defs.ccT3), and ORT builds a session for every one of them — I probed all three against a plain-ORT session and all three loaded. Our kernel implements packeduint8only. Fixing defect 1 by advertising the union therefore claimed a float16-zero-point node and then failed insideexecutewithzero_points must have dtype Uint8— turning a model that previously produced an answer into a hard error.KernelRegistryEntrygainsinput_dtype_constraints: &'static [(usize, &'static [DataType])]. A listed slot is checked against its own set instead of the union; unlisted and absent slots keep the union, which is exact for every uniform op. Populated straight from the kernels'require_dtypecalls:MatMulNBitsA, 2scales, 5biasB, 3zero_pointsg_idxQLinearMatMula_scale, 4b_scale, 6y_scaleThe
QLinearMatMulhalf matters for the same reason: its scales flow throughARITH_DTYPES, which admits f16/bf16, and opset 21 spec-allows those — so defect 3's fix would newly have claimed a configqlinear_matmul.rs:316rejects.Remaining gap, stated plainly
Float / float16 / bfloat16
zero_pointsonMatMulNBitsis a quantization semantic this EP does not implement. It is declined atGetCapability, anda_float_zero_point_matmul_nbits_is_declined_not_claimed_and_failedpins that it is declined and still computed rather than claimed and failed. That is a genuine coverage gap in the matmul family, recorded here rather than hidden: closing it means implementing the unpacked float zero-point dequant path, which is a new numeric semantic and does not belong in an urgent claim fix.Falsifiers
no_matmul_family_node_escapes_to_the_ort_cpu_epbuilds ten models from scratch, runs each under a session withsession.disable_cpu_ep_fallback=1, asserts via ORT's ownrecord_ep_graph_assignment_infothat the op is oncpu_epand thatops_not_on_our_ep()is empty, then compares elementwise against a second session with no EP appended — i.e. pure ORT CPU. A non-degeneracy guard requires at least half the outputs to be non-zero and finite first, so a both-sides-zero comparison cannot pass vacuously.matmul_f32_decodematmul_f32_prefillmatmul_f16_decodegemm_f16_decodenbits4_decodenbits8_prefillnbits4_block512nbits4_f16_decodeqlinear_u8qlinear_i8Two further cases outside the sweep:
mul_self(defect 4) andnbits_float_zp(the declined-not-crashed guard).nbits4_block512is the negative control: ORT rejects block_size 512 at kernel construction (matmul_nbits.cc:131), we run it, and the test asserts ORT refuses so the claim in the assignment doc cannot go stale silently.nbits4_f16_decodeexists because fixing defect 1 widens what we claim — f16 activations were unreachable before, and this pins that they are now correct rather than merely accepted.a_node_that_names_one_value_twice_is_bound_onceis the direct falsifier for defect 4:Mul(X, X)on a graph input, no initializers, no constant-sharing pass involved. It asserts the node lands on our EP and that every element equalsx*xexactly.Each fix falsified independently
Reverting one source file at a time on this branch, with the test unchanged:
nbits4_decode— session creation fails,ours=[],others=["MatMulNBits"]compute.rsonly (shape rules)nbits4_decode—MatMulNBits: Y must have shape [1, 4096], got [4096, 1, 16]ep.rsonly (slot dedup)qlinear_u8—input slot 7 maps to ORT input 7, but ORT bound only 7 input(s)mul_self—input slot 1 maps to ORT input 1, but ORT bound only 1 input(s)Plus seven hardware-independent unit tests: four in
compute.rspinning the two new shape rules — including an explicitassert_eq!(wrong, vec![vec![4096, 1, 16]])so the plain-MatMulalias cannot return unnoticed — two inkernels/mod.rsasserting the quantized edge dtypes are advertised, and one inep.rsasserting the per-slot filter declines a float16-zero-point node and a uint8-activation node while accepting the packed-uint8 one.Performance
No performance claim. This PR changes nothing about the kernels; it changes which kernels run at all. It does not make anything faster, and the ratios in
CPU_MATMUL_ASSIGNMENT.mdare not restated here — see below.What this does to the published matrix
The
MatMulNBitsandQLinearMatMulrows inCPU_MATMUL_ASSIGNMENT.mdwere measured withbench_generic, which drives the kernels natively, so those numbers remain valid as kernel-level measurements. What is now known to be wrong is the implication that the plugin path ever reached those kernels. Re-stating the matrix as plugin-path numbers is deliberately not in this PR — it is a doc change that depends on #1097's rewrite of the same file, and folding it in here would make an urgent correctness fix conflict with an open architectural PR.Related, not fixed here
Same class of bug, outside this PR's remit, reported to the owners rather than fixed silently:
("MoE" | "QMoE", "com.microsoft") => F32_ONLY— both carryuint8weight edges.("GatherBlockQuantized", "com.microsoft") => FLOAT_DTYPES— same.Both are very likely dropped by the dtype filter for the same reason as defect 1.
Verification
cargo test -p onnx-runtime-ep-cpu-plugin --test plugin_ort_e2e— 50 passed (was 47; +3)cargo test -p onnx-runtime-ep-plugin --lib— 229 passed (was 224; +5)cargo test -p onnx-runtime-ep-cpu --lib— 1317 passed; with--features mlas1334 passed (+2 dtype tests)cargo test -p onnx-runtime-ep-shared-mock-plugin— passes (registry-entry field addition)cargo fmt --all -- --checkcleancargo clippy --all-targetsclean both with and without--features mlasHost: AMD EPYC 9V74, 32 vCPU, AVX2+FMA+F16C, ORT 1.27.0.
Independent review
Reviewed by an independent Opus reviewer against a separate worktree, with the ORT 1.27.0 source for
partitioning_utils.ccpulled at the prebuilt's own commit. Verdict was REQUEST CHANGES on the over-claim above; all findings are addressed in905f7c47c:MatMulNBitsclaimed then fails at RunQLinearMatMulclaimed then fails at RunQLinearMatMulslot tableqlinear_u8— but a directMul(X, X)case was added anyway, since that one covers a graph input rather than an initializerRound two returned APPROVE. The reviewer disabled
input_dtype_constraints_for_opforMatMulNBitsand confirmed the float-zero-point test then fails, i.e. the per-slot list is load-bearing and nothing else catches it; wrote an independent opset-21 f16-scaleQLinearMatMulprobe and confirmed it is declined and still computed; reverted the slot mapping to positional and confirmed bothmul_selfandqlinear_u8fail; and checked all fourteen constrained slots against the two kernels'require_dtypecalls. One useful correction: forQLinearMatMulthe per-slot list is redundant with a pre-existing nativeunsupported_reasongate, so it is defence in depth there and load-bearing only forMatMulNBits.The reviewer independently confirmed from ORT's
MakeComputeCapabilitythat the fused-node input set is deduplicated byNodeArgidentity in first-appearance order, over graph inputs and initializers alike, and that the multi-node routing path already keys on value position so it stays consistent with the fast path.Post-review fix
f8f7a486a— the generated-model element-type constants were annotatedu32.ONNXTensorElementDataTypeis ac_uint, which bindgen renders asu32on Linux buti32on Windows, soRust (Windows ARM64)andRust coverage (Windows x86_64)failed to compile the test file while every Linux lane was green. They now use ORT's own alias. Test-only, no behaviour change.