Dispatch-overhead stack: benchmark grid, 7 -> 4 FFI calls per Run, routed-loop fixes (2.31 -> 1.71 at depth 100) by justinchuby · Pull Request #1397 · justinchuby/onnx-genai

justinchuby · 2026-08-19T04:32:06Z

Note on scope. Enabling auto-merge across the stack caused each stacked PR to merge into its parent feature branch (those branches are unprotected; only main requires checks), and GitHub retargeted the children. No main protection was bypassed and no work was lost, but this PR now carries the whole dispatch-overhead stack below #1387, not just the memory-info change:

Add the #1077 dispatch-overhead benchmark grid #1394 — the Reduce cpu ep kernel dispatch overhead #1077 benchmark grid

this PR — device_mem_info, 7 → 4 FFI calls per Run

Instrument the routed loop, and fix the O(depth^2) retirement scan (2.31 -> 1.95 at depth 100) #1401 — routed-loop instrumentation, the O(depth²) retirement scan, lazy absent strides, and the 8-element dispatch-isolating cases

Each was independently reviewed before it merged; the sections below cover the memory-info change, and #1394/#1401 retain their own descriptions and review records. #1387 (probe + lazy formatting) remains a separate PR to main.

Stop asking ORT three times where the inputs live

Part of #1077. Stacked on #1394 (which measured the cost this removes).

The finding

Every Compute call resolves where a routed subgraph's intermediate tensors should be allocated. On a host EP — and the CPU plugin is always a host EP — that resolution cost 7 FFI calls for a single-input node to arrive at the answer "host memory". Three of them were asking ORT for something already in hand.

device_mem_info scans the inputs looking for a device-resident one:

per input: GetValue → GetTensorMemoryInfo → mem_info_is_device
on a host EP it never finds one, so the scan runs to completion
it then falls through to "use input 0's memory info" — and fetches input 0 again, two more calls for a value it had already looked at and discarded
the caller then asked mem_info_is_device on the result — a third query for a fact the function had just established internally and thrown away

The change

The scan remembers input 0's memory info as it walks, and returns (mem_info, is_device) so the caller does not have to re-derive the second half.

The fallback fetch stays, for the one case that genuinely needs it: no input count available, so the scan never ran and nothing was remembered.

Why this is behaviour-identical

The scan visits the same inputs, in the same order, and returns on the same condition.
The remembered value is exactly what the re-fetch would have produced: both are input 0's OrtMemoryInfo, and it is immutable for the lifetime of the value.
The recon_mem_info branch reports is_device = true without asking. That branch is only reachable when the scan found a device-resident input, so it is device-resident by construction — which is precisely what the removed call was confirming.

Evidence

mem_info_cost (new, in compute.rs) drives device_mem_info through a hand-built OrtApi and pins the call count.

scenario	before	after
host EP, 1 input	7	4
device-resident input (short-circuit)	3	3
no input count (fallback path)	2	2
zero inputs	0	0

The "before" column is not a recollection — it is the same test run against the previous implementation.

Three FFI calls per Run, off the fixed per-Run cost that #1394 measures at ~0.9 µs worse than ORT's.

Why counters and not a stopwatch

This machine has not been quiet enough all session for a trustworthy wall-clock A/B (mpstat shows cores at 75–78% idle against a ≥93% gate). An FFI-call count is deterministic and load-independent, so it is reproducible now and will still be reproducible on someone else's machine. That is the case for #1387 existing.

Also pinned

The three non-host paths, so a future change cannot quietly break them while the host count stays pretty: device-resident short-circuit, missing input count, and a node with no inputs (reports no memory info rather than inventing one).

Validation

cargo test -p onnx-runtime-ep-plugin --lib — 253 pass (debug; --release silently disables debug_assert!)
same with --features dispatch_probe — 264 pass
plugin_ort_e2e — 55/55, including every_assigned_node_is_also_executed_by_this_ep
clippy clean in both feature configurations

No ORT CPU fallback. No MLAS.

Independent review (Opus 4.8, adversarial) — APPROVE WITH NITS, all findings addressed

Commissioned specifically to attack "behaviour is unchanged in every case". It held for every EP that ships, with one crack worth fixing.

Finding 1 (MINOR) — the recon branch's is_device = true was convention, not construction.

Upheld, and fixed properly. My comment claimed the reconstruction is "device-resident by construction". The reviewer traced it: recon_mem_info is built from the EP's reported device_type, and device_staging is attached whenever an EP supplies a HostToDeviceCopier — and nothing in that trait requires a copier-providing EP to be non-CPU. Every one today is (CUDA→GPU, QNN→NPU), which is exactly what made the hardcode look safe. A CPU-typed EP that grew a copier would build a host-typed memory info and the hardcode would hand it to the device scratch path — host pointers treated as device pointers, the failure this code exists to prevent, inverted. The old code had no such exposure because it asked ORT.

Rather than paper over it with a debug_assert, the invariant is now structural. mem_info_is_device is exactly device_type != CPU, and the reconstruction is created by passing that same device_type to CreateMemoryInfo_V2 — so recording it at construction reproduces the old query in every case, including the one that does not exist yet, and still costs no FFI call.

Finding 2 (MINOR) — the one branch whose reporting changed was untested. Correct: all four original tests passed staging = None and never reached it. Three tests added — CPU-typed reconstruction is not reported as a device, device-typed one is (4 calls, not 5), and a real device input still wins over the fallback.

Falsified: restoring the hardcoded true fails exactly one test, the one naming the case; the other six correctly stay silent.

Finding 3 (NIT) — test 2's == 4 does not distinguish new from old. Accurate. The device-input path cost 4 before as well; the real guard in that test is the tuple value, not the count. Left as-is (it pins the device path against future regressions), noted here so nobody reads it as evidence for this PR.

Claims the reviewer tried to break and could not

These are now load-bearing rather than merely unchallenged:

Pointer stability — the remembered input-0 memory info cannot differ from what the re-fetch produced: same OrtValue, straight-line code with no intervening mutation, and the pointer's lifetime window is identical in both versions.
The fallback path does not hardcode anything — it still queries, so a device-resident input on that path still reports true, exactly as before.
Zero-input nodes — identical (None, not an invented memory info).
Caller equivalence — Option<(*const _, bool)> is Copy; no move, no shadowing, no ordering change.
mem::zeroed::<OrtApi>() — sound; all fields are Option<extern fn>, zero is the null-fn niche, and the sentinel pointers are never dereferenced.
The ffi_coverage tripwire stays honest — 9 members / 9 ort_call() sites, unchanged and correct: this PR removes dynamic calls, not static sites.
Concurrency — entirely stack-local; no new shared mutable state.
No vacuous test — the reviewer was briefed that I have twice shipped a test that passed against the broken version, and specifically hunted a third. None found.

Two allocations sat on the per-Run dispatch path purely to serve diagnostics that are off. `staging_log` returns immediately when the transfer trace is disabled, but its argument is evaluated first: `staging_log(&format!(…))` builds a `String` and formats every interpolated value on the way to a function that drops it. One of these is in `compute_execute`, before any real work, so every dispatch in every production run paid for it. A `staging_log!` macro checks the gate before formatting, which makes a disabled trace cost one relaxed `OnceLock` load. Thirteen call sites convert; the emitted text is unchanged, verified by running the e2e suite with `ONNX_GENAI_PLUGIN_TRANSFER_TRACE=1` and confirming the `compute_execute enter` line still appears. `read_inputs` built `format!("input {i}")` for every input on every Run to label an error that is almost never raised. `validate_dims` now takes `impl Display`, so the call site passes `format_args!`, which borrows its arguments and formats only if a message is actually produced. Same for `transfer.rs`'s `CopyTensors[{i}]`. No side-effecting expression moved inside a now-conditional macro: the only mutation near a converted site (`to_stage.push`) is outside it. The pinned counts in `dispatch_cost` drop from 5 allocations per input to 4, and from `1 + 3*4` to `1 + 3*3` for three inputs. That is the point of #1387 — the improvement is a number changing in a test, reproducible on any machine, rather than a claim resting on a benchmark run on a shared box. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Issue #1077 asks whether our per-Run dispatch overhead is worse than ORT's. Every measurement so far has been a single-node elementwise A/B, which reports one number in which two very different costs are indistinguishable: what ORT charges once per Run for session bookkeeping and feed/fetch marshalling, and what we charge per node for dispatch. Improving the second while reporting the sum makes progress look smaller than it is, and hides which one is actually the problem. This adds a grid whose axes separate them. `chain_case` builds a chain of `depth` identical nodes, so depth 1, 10 and 100 read the per-node slope with the per-Run constant divided out. `Identity` gives the floor — the kernel is a memcpy, so essentially everything the timer sees is the cost of getting there and back. `small_matmul_case` checks the answer is not an artefact of trivial kernels, with both operands as runtime inputs so neither side prepacks and the weight stays out of the fixture. The `_dyn` rows repeat the static ones with a symbolic batch dimension, so a fast path that only works when shapes are known at session build cannot pass unnoticed. Gated behind `NXRT_MM_BENCH_GRID`: `1` appends the grid to the existing cases, `only` runs it alone, which is what a dispatch measurement wants — nothing else resident and nothing else warming caches. Default behaviour is unchanged. First run (1 thread, 15 iters, interleaved, on a machine that was NOT quiet — treat the absolute numbers as provisional, but the slope is a within-arm quantity and survives the noise): relu_1 ours 3.5us ort 2.2us relu_10 ours 12.7us ort 8.5us relu_100 ours 107us ort 64.6us per-node: ours 1.03us, ORT 0.65us -> we are ~0.4us/node behind per-Run: ours ~2.5us, ORT ~1.6us -> and ~0.9us behind before any node runs That is the decomposition #1077 was missing. It also turned up something unrelated to dispatch: `grid_matmul_512x512` is 2.74x, far above dispatch scale (77us vs 28us), so f32 MatMul is a kernel-side gap and needs its own work rather than being read as dispatch overhead. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Before any kernel runs, every Compute call resolves where a routed subgraph's intermediates should be allocated. On a host EP — the CPU plugin, always — that resolution took seven FFI calls for a one-input node to arrive at the answer "host memory", and three of them were asking for something already known. `device_mem_info` scans the inputs for a device-resident one. On a host EP it never finds one, so the scan runs to completion and falls through to "use input 0's memory info" — which it fetches again, two more calls for a value it looked at first and discarded. The caller then asked whether the result was device memory, a third query for a fact the function had just established and thrown away. The scan now remembers input 0 as it goes and returns the device-ness it already computed. The fallback fetch remains for the one case that genuinely needs it: no input count available, so the scan never ran. Behaviour is unchanged in every case. The scan visits the same inputs in the same order and returns on the same condition; the reconstructed device memory info is device-resident by construction, which is what the call it replaces was confirming. Measured, not asserted: the `mem_info_cost` tests drive `device_mem_info` through a hand-built `OrtApi` and pin the count at 4. Running that same test against the previous implementation reports 7. Three FFI calls per `Run`, on the fixed per-`Run` cost that #1394 measures as ~0.9us worse than ORT's. This is the kind of evidence #1387 exists to produce — a number that is reproducible on a loaded machine, which matters here because the box has not been quiet enough for a trustworthy wall-clock A/B all session. Also pinned: a device-resident input short-circuits and reports itself, a missing input count still falls back to input 0, and a node with no inputs reports no memory info rather than inventing one. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Review of the parent commit accepted that the rewrite is behaviour- identical for every EP that exists, but found the argument for one line weaker than the comment claimed. The fallback returned `(recon.ptr, true)` on the grounds that the reconstruction is device-resident by construction. It is not, quite. `recon_mem_info` is built from the `device_type` the EP reports, and `device_staging` is attached whenever an EP supplies a host-to-device copier. Nothing in `HostToDeviceCopier` requires a copier-providing EP to be non-CPU. Every one today is -- CUDA reports GPU, QNN reports NPU -- which is what makes the hardcode look safe. A CPU-typed EP that grew a copier would build a host-typed memory info, and the hardcode would hand it to the device scratch path as though it were device memory: host pointers treated as device pointers, which is the failure the surrounding code exists to prevent, inverted. The old code did not have this exposure, because it asked ORT. `mem_info_is_device` is exactly `device_type != CPU`, and the reconstruction is created by passing that same `device_type` to `CreateMemoryInfo_V2`. So recording it at construction reproduces the old query in every case, including the one that does not exist yet, and still costs no FFI call. The invariant is now structural rather than conventional, which is the difference between a comment and a guarantee. Three tests cover the branch, which the previous commit left untested -- every existing test passed `staging = None` and never reached it. A CPU-typed reconstruction is not reported as a device; a device-typed one is, for 4 FFI calls rather than 5; and a real device input still wins over the fallback. Falsified: restoring the hardcoded `true` fails exactly one test, the one naming the case, and the other six correctly stay silent. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

….31 -> 1.95 at depth 100) (#1401) ## The per-node gap is not dispatch overhead — and it was uninstrumented Part of #1077. Stacked on #1397. ### 1. A ten-node chain is one dispatch, not ten I had been reading #1394's per-node cost as ORT callback overhead. It is not. `query_capabilities_filtered` unions adjacent supported nodes into one convex partition, so a Relu chain is claimed as **a single fused subgraph** and ORT calls `Compute` **once**. Confirmed by running the chain with the staging trace on: ``` [plugin/staging #982] compute_execute enter: entries=10 routing=true device_staging=false ``` emitted three times for three `Run`s — one callback carrying ten nodes. So the per-node gap contains **no FFI at all**. ORT pays real per-kernel dispatch ten times and still beats us; we pay it once and lose anyway. This also exposed a hole in #1387: `TensorBind`, `DispatchLookup`, `KernelInvoke` and `NodeExecuted` were all on the **single-node fast path** — the path a multi-node graph never takes. A probe that misses the interesting path will confirm whatever you already believed. The routed loop now reports the same phases per node. ### 2. The grid was measuring memory bandwidth Every chain ran at 4096 f32. A Relu node at that width moves 32 KiB, ~0.6 µs on a quiet machine — comparable to the entire per-node cost I was attributing to dispatch. The tell was in the data and I had missed it: **every case came back at almost exactly 1.5×**, from a 3 µs Identity to a 100 µs Relu chain. A uniform ratio across a 30× range of work is a *multiplicative* effect; dispatch overhead is *additive*. Added 8-element cases. At that width the kernel is one masked vector op and what remains is overhead. The two widths bracket the question instead of conflating it — TINY is the overhead floor, 4096 the throughput regime. The isolated gap is **larger** than the one I had been reporting: | depth | ours | ORT | per-node ours | per-node ORT | |---|---|---|---|---| | 1 | 2.6 µs | 1.7 µs | — | — | | 10 | 6.6 µs | 4.0 µs | 0.44 µs | 0.26 µs | | 100 | 54.2 µs | 23.5 µs | **0.52 µs** | 0.22 µs | ORT holds ~0.22 µs/node from 10 to 100. Ours *rose* from 0.44 to 0.52, so the ratio grew with depth — 1.52 at one node, 2.31 at a hundred. That is the opposite of a fixed per-`Run` cost being amortised. ### 3. A per-node cost that depends on node count is not a per-node cost At each node the routed loop walked the **whole** last-reader table to find the buffers that node was the final reader of. A chain allocates a buffer per node, so that is **O(depth²)** — at depth 100, 9,900 comparisons to perform 99 retirements. Inverting the map once per `Run` into a CSR index makes each node touch only its own retirements. | depth 100, 8-wide chain | before | after | |---|---|---| | ours p50 | 54.2 µs | **46.0 µs** | | ORT p50 (control) | 23.5 µs | 23.6 µs | | ratio | 2.310 | **1.951** | | per-node ours | 0.52 µs | **0.44 µs** | Per-node now **agrees between depth 10 and depth 100**, which is the result that says the quadratic term was the whole of the difference rather than part of it. ORT moved 0.4%, so the machine did not change underneath me. ### 4. Allocations removed from the per-node path - `node_label` built a `String` for every node of every `Run` to interpolate into error messages that are not produced — same shape as the eager `format!`s #1392 removed from `read_inputs`. `prepare_workspace` and `workspace_trace_line` take `impl Display`; the call site passes `format_args!`. - `absent_shapes` deep-copied `output_shapes` (a `Vec<Vec<usize>>`, so one allocation plus one per output) and then only ever read it. It borrows now. Both paths had it. - **Absent-output strides** were built for *every* output of every node and read only for slots that turned out to be absent — an allocation per output, per node, discarded unused. Absent outputs are the exception, so on the ordinary path all of it was thrown away. The obvious fix — skip it when nothing is absent — leans on "`absent_scratch` is empty iff no slot is `Absent`", an invariant with almost no coverage (one lib test mentions absent slots; no e2e case exercises one). Optional slots are something #1077 explicitly must not break, so resting correctness on an untested invariant is the wrong trade for two allocations. The storage instead holds one entry per absent slot, keyed by the index `Absent(idx)` already carries, so the keys agree **by construction** and a node with no absent outputs collects an empty iterator (which does not allocate). Keying is what gets tested, because it is invisible in the common case: with one absent output at slot 0 the two indices coincide and any mistake hides. The test uses absent slots that are neither first nor contiguous, with a different shape at each. Falsified: re-keying by position fails it. ### 5. Cumulative effect | depth 100, 8-wide chain | baseline | + CSR retirement | + lazy absent strides | |---|---|---|---| | ours p50 | 54.2 µs | 46.0 µs | **41.0 µs** | | ORT p50 (control) | 23.5 µs | 23.6 µs | 23.9 µs | | **ratio** | 2.310 | 1.951 | **1.713** | | per-node ours | 0.52 µs | 0.44 µs | **0.39 µs** | **24% faster at depth 100**, with the ORT arm drifting 1.7% across the three runs as the noise control. Depth 1 is unchanged (1.516) — as expected, since none of this is on the fixed per-`Run` cost. ### Correctness The retirement rewrite is checked **against the scan it replaces**, not against hand-written expectations: for a chain, a twice-read buffer, unread and out-of-range buffers, a node retiring three buffers at once, and a subgraph with no buffers, the index must yield exactly what the scan yielded, in the same ascending order. Plus: every buffer retires exactly once — retiring twice would recycle storage that is still live. ### Validation - `cargo test -p onnx-runtime-ep-plugin --lib` — 258 pass (debug; `--release` silently disables `debug_assert!`) - `--features dispatch_probe` — clean - `plugin_ort_e2e` — **55/55**, including `every_assigned_node_is_also_executed_by_this_ep` - clippy clean in both feature configurations - 6000 interleaved iterations, pinned to two cores, 7 of 8 sampled cores >95% idle ### What remains 0.39 µs/node against ORT's 0.22, and 2.5 µs vs 1.7 µs fixed per `Run`. The per-node remainder is now dominated by the ~8–10 heap allocations each node still makes (`kernel_inputs`, `infer_shapes`' `Vec<Vec<usize>>`, `absent_strides_storage`, `all_output_views`, `node_ort_operands`, `prepare_workspace`'s metadata vector). Thread-local reusable scratch is the route; per-session caching is barred by ORT's concurrent-`Run` hazard. Tracked in #1077. No ORT CPU fallback. No MLAS. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

github-actions · 2026-08-19T11:51:41Z

🔴 Benchmark Regression Detected

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_f16_threads=8/32x1024x1024`	90.13 µs	123.67 µs	+37.2%
⚠️	`gather/medium_f16_threads=1-internal/32768`	2.53 µs	3.09 µs	+22.3%
⚠️	`add/large_f16_threads=1-internal/4194304`	1.88 ms	2.21 ms	+17.8%
✅	`gather/small_bf16_threads=1-internal/4096`	516.6 ns	583.5 ns	+12.9%
✅	`matmul/medium_generic_bf16_threads=8/32x512x512`	465.46 µs	518.21 µs	+11.3%
✅	`matmul/large_generic_f16_threads=1/32x1024x1024`	89.24 µs	95.11 µs	+6.6%
✅	`matmul/small_generic_bf16_threads=1/1x256x256`	31.00 µs	32.88 µs	+6.0%
✅	`add/large_bf16_threads=1-internal/4194304`	1.90 ms	2.01 ms	+5.4%
✅	`matmul/large_generic_f32_threads=1/32x1024x1024`	9.53 ms	10.05 ms	+5.4%
✅	`matmul/medium_generic_f16_threads=8/32x512x512`	36.58 µs	38.49 µs	+5.2%
✅	`reduce_mean/large_f32_threads=1-internal/262144`	1.13 ms	1.18 ms	+4.0%
✅	`add/medium_f32_threads=1-internal/262144`	28.47 µs	29.03 µs	+2.0%
✅	`block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024`	123.53 µs	125.95 µs	+2.0%
✅	`gather/small_f16_threads=1-internal/4096`	531.5 ns	541.7 ns	+1.9%
✅	`block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1`	599.23 µs	608.77 µs	+1.6%
✅	`qwen3_sampling_processors/top_k_full_sort_baseline`	2.37 ms	2.36 ms	-0.4%
✅	`matmul/medium_generic_bf16_threads=1/32x512x512`	555.16 µs	552.43 µs	-0.5%
✅	`add/small_f32_threads=1-internal/1024`	251.4 ns	250.1 ns	-0.5%
✅	`block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1`	136.07 µs	133.34 µs	-2.0%
✅	`matmul/large_generic_f32_threads=8/32x1024x1024`	4.65 ms	4.49 ms	-3.4%
✅	`qwen3_sampling_processors/top_k_top_p_full_sort_baseline`	6.34 ms	6.11 ms	-3.7%
✅	`kv_cache/alloc_dealloc_pages`	42.22 µs	40.18 µs	-4.8%
✅	`reduce_mean/medium_f32_threads=1-internal/65536`	309.58 µs	294.41 µs	-4.9%
✅	`sampling_latency/min_p_per_token`	218.15 µs	207.40 µs	-4.9%
✅	`add/small_bf16_threads=1-internal/1024`	499.2 ns	471.5 ns	-5.6%
✅	`reduce_mean/small_f32_threads=1-internal/4096`	18.58 µs	17.45 µs	-6.1%
✅	`matmul/small_generic_f32_threads=1/1x256x256`	41.39 µs	38.85 µs	-6.1%
✅	`block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024`	907.10 µs	847.66 µs	-6.6%
✅	`matmul/medium_generic_f16_threads=1/32x512x512`	35.39 µs	33.05 µs	-6.6%
✅	`gather/small_f32_threads=1-internal/4096`	813.6 ns	757.2 ns	-6.9%
✅	`matmul/large_generic_bf16_threads=8/32x1024x1024`	1.59 ms	1.47 ms	-7.1%
✅	`qwen3_sampling_processors/top_k_top_p_fast`	715.71 µs	663.76 µs	-7.3%
✅	`gather/large_f32_threads=1-internal/131072`	53.84 µs	49.93 µs	-7.3%
✅	`matmul/small_generic_f16_threads=8/1x256x256`	37.06 µs	34.05 µs	-8.1%
✅	`logit_processing/seven_processor_chain_per_step`	357.71 µs	327.25 µs	-8.5%
✅	`sampling_latency/top_p_per_token`	455.67 µs	415.33 µs	-8.9%
✅	`matmul/small_generic_f16_threads=1/1x256x256`	35.97 µs	32.74 µs	-9.0%
✅	`qwen3_sampling_processors/top_p_fast_after_top_k`	594.67 µs	539.22 µs	-9.3%
✅	`gather/large_bf16_threads=1-internal/131072`	17.16 µs	15.36 µs	-10.5%
✅	`gather/medium_f32_threads=1-internal/32768`	5.63 µs	5.04 µs	-10.6%
✅	`add/small_f16_threads=1-internal/1024`	650.5 ns	578.5 ns	-11.1%
✅	`matmul/large_generic_bf16_threads=1/32x1024x1024`	2.33 ms	2.06 ms	-11.6%
✅	`qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline`	4.44 ms	3.92 ms	-11.6%
✅	`tokenization/decode_tokens_per_second`	7.36 ms	6.48 ms	-12.1%
✅	`matmul/small_generic_bf16_threads=8/1x256x256`	41.33 µs	35.61 µs	-13.8%
✅	`sampling_latency/greedy_per_token`	3.69 µs	3.15 µs	-14.6%
🟢	`grammar_masking/llguidance_compute_mask/32`	92.25 µs	77.72 µs	-15.8%
🟢	`matmul/medium_generic_f32_threads=1/32x512x512`	2.96 ms	2.47 ms	-16.5%
🟢	`qwen3_sampling_processors/top_k_partial_selection`	185.12 µs	148.21 µs	-19.9%
🟢	`tokenization/encode_tokens_per_second`	509.78 µs	405.93 µs	-20.4%
🟢	`matmul/medium_generic_f32_threads=8/32x512x512`	1.60 ms	1.27 ms	-20.5%
🟢	`sampling_latency/top_k_per_token`	68.18 µs	52.77 µs	-22.6%
🟢	`gather/medium_bf16_threads=1-internal/32768`	3.59 µs	2.75 µs	-23.4%
🟢	`matmul/small_generic_f32_threads=8/1x256x256`	51.44 µs	37.81 µs	-26.5%
🟢	`add/large_f32_threads=1-internal/4194304`	1.12 ms	814.88 µs	-27.0%
🟢	`add/medium_bf16_threads=1-internal/262144`	192.08 µs	139.78 µs	-27.2%
🟢	`add/medium_f16_threads=1-internal/262144`	198.93 µs	138.40 µs	-30.4%
🟢	`gather/large_f16_threads=1-internal/131072`	20.04 µs	13.69 µs	-31.7%
🟢	`block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024`	95.70 µs	55.29 µs	-42.2%

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: { 2.49 3.07 5.51 }

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)

…d (2.31 -> 1.71 at depth 100) (#1387) > **Scope note.** This PR now carries the whole #1077 dispatch-overhead stack. Enabling auto-merge across the stacked PRs caused each to merge into its *parent feature branch* — those branches are unprotected; only `main` requires checks — and GitHub retargeted the children, collapsing the stack into this one. **No `main` protection was bypassed and no work was lost.** This PR still waits on required CI before it can reach `main`. **What is in here, in dependency order — each was independently reviewed as its own PR before it merged:** | was | what | evidence | |---|---|---| | #1387 | deterministic phase/counter instrumentation for the dispatch path | Opus 4.8 REQUEST CHANGES → all findings fixed | | #1392 | stop building debug strings production throws away | 5 → 4 allocations per input per `Run` | | #1394 | the #1077 benchmark grid (depth 1/10/100, static + dynamic) | produced the first real decomposition | | #1397 | stop asking ORT three times where the inputs live | **7 → 4** FFI calls per `Run`, Opus 4.8 APPROVE WITH NITS → fixed | | #1401 | routed-loop instrumentation; O(depth²) retirement scan; lazy absent strides; 8-element dispatch-isolating cases | **2.310 → 1.713** at depth 100 (−24%) | **Headline result** — 8-wide Relu chain, depth 100, 6000 interleaved iterations, pinned cores, ORT arm as drift control (moved 1.7%): | | baseline | after | |---|---|---| | ours p50 | 54.2 µs | **41.0 µs** | | ratio vs ORT | 2.310 | **1.713** | | per-node | 0.52 µs | **0.39 µs** | Two of my own earlier conclusions were falsified along the way and are documented in #1077: a 10-node chain is **one** fused dispatch (so the per-node gap contains no FFI at all), and the original grid was measuring **memory bandwidth**, not dispatch. --- ## Why Issue #1077 says our per-`Run` dispatch overhead is higher than ORT's. Every time I have had to answer *where* that overhead is, the method has been the same: hand-edit `Instant` probes into the hot path, run the A/B, read the numbers, tear the probes out. That is slow, it is unreviewable, and it leaves nothing behind — the next person starts from zero, and nothing catches a regression that quietly adds a round trip back. This makes the measurement a permanent, reviewable part of the crate. ## What `dispatch_probe`, gated behind a **non-default** Cargo feature, wired through the CPU dispatch path: | Phase | Where | |---|---| | `CallbackEntry` | `compute_execute` entry, up to the first real work | | `MetadataQuery` | `read_inputs` | | `TensorBind` | single-node input view construction | | `Allocate` | `allocate_output` | | `DispatchLookup` | `infer_shapes`, `prepare_workspace` | | `KernelInvoke` | `execute_with_workspace` | | `StatusCrossing` | `status_with_code` | Plus event tallies for the three things actually worth removing: `OrtFfiCall`, `DispatchAlloc`, `StatusCreated` (and `ComputeExecute` / `NodeExecuted` as denominators). ## Design notes **The counters are kept twice, deliberately.** The thread-local copy is the precise one — a dispatch is a single-threaded story, and only per-thread accumulation lets one measurement be read without a concurrently-running sibling test bleeding into it. That isolation is what makes the exact-count assertions below possible at all; I tried global-only first and the counts were nondeterministic under the default parallel test runner. The global mirror exists because ORT chooses the thread that runs `Compute`, and it need not be the thread that called `Run`. The e2e harness loads this EP as a cdylib through `RegisterExecutionProviderLibrary`, so it reads the counters back through an exported `nxrt_dispatch_probe_snapshot` symbol — and it has no way to be on ORT's worker thread. A thread-local-only probe would report zero there, which reads as "we made no FFI calls" rather than "you are looking at the wrong thread". Writing both costs one `Cell` store and one relaxed `fetch_add`. **Production pays nothing.** Without the feature, `PhaseGuard` is a zero-sized type with no `Drop` impl and every entry point is an `#[inline(always)]` empty function. Call sites use `guard.end()` rather than `drop(guard)` so the production build does not read as dropping a `Drop`-less ZST. **Timing is gated a second time**, on `ONNX_GENAI_PROFILE_DISPATCH=1`. A test asserting FFI call counts wants the counters but emphatically does not want two `Instant::now()` calls added to every phase it is measuring. ## The counts are pinned, not documented The point of the module is the regression guard, so the numbers are asserted rather than written in a comment. `kernel_ctx` gains a hand-built `OrtApi` — zeroed, then filled in with only the entry points `read_inputs` actually uses, so any call it makes that the test did not anticipate faults loudly instead of silently succeeding — which lets `read_inputs` run with no live ONNX Runtime at all. Today, per `Run`: - **8 FFI round trips** for one input (within `read_inputs`) — one shared `GetInputCount`, then 7 per input - **5 heap allocations** — the `Vec<OwnedInput>` once, then `dims`, the `format!("input {i}")` label, `shape` and `strides` per input - three inputs cost `1 + 3×7` and `1 + 3×4`, so this is a linear model rather than a single data point - an absent optional input short-circuits at 2 calls, which pins the optional-slot path Any change that adds a round trip or an allocation to the per-`Run` path now has to come here and edit a number, in a diff a reviewer can see. ## Verification Both feature configurations: `cargo fmt`, `cargo clippy --all-targets` (zero warnings), and the full `--lib` suite — 251 tests without the feature, 257 with it. The real ORT-driven suite passes unchanged: **55/55 `plugin_ort_e2e`** with `NXRT_REQUIRE_ORT_TESTS=1`, including `every_assigned_node_is_also_executed_by_this_ep`. **Every test here was verified by mutation**, because a test that passes against the known-broken version is not a test: - Injecting one extra `get_type_shape` call and one extra `Vec` allocation into `read_inputs` failed exactly the three count assertions and correctly left the other two green. - The zero-cost assertion **was vacuous when I first wrote it** — it compared `snapshot()` after counting against `snapshot()` before, which are two calls to the same constant function in a disabled build. It passed unchanged against a mutant whose "disabled" `count_n` recorded into a static and whose `snapshot` returned non-zero. It now compares against `Counters::default()` — an absolute claim rather than a self-consistency one — and fails against that mutant. It also asserts `!needs_drop::<PhaseGuard>()`, closing the one way the production guard could regain a cost while staying zero-sized. ## Review round 1 — Opus 4.8, `REQUEST CHANGES`, all addressed The review found a real defect, and it was the module's central claim. **Blocker: `OrtFfiCall` was documented as a per-`Run` total but only `read_inputs` and `allocate_output` were instrumented.** On the happy path of a single-node `Run`, `device_mem_info`, `ort_input_mem_info`, `mem_info_is_device`, `alloc_scratch`, `host_pool::install` and `CreateStatus` all reach into ORT and none were counted. A probe that under-counts while calling itself exact is worse than no probe — it invites exactly the wrong conclusion about where the overhead is. Fixed by instrumenting every leaf `OrtApi` call in the crate, and then by making that property **self-enforcing** rather than a promise. The new `ffi_coverage` tests scan the source of each file that touches `OrtApi` and extract the API members it names — ORT's C API is `CamelCase` and our own fields are `snake_case`, so `.CamelCase` isolates them cleanly — then fail if that set or the instrumentation count moves. Verified by mutation: adding a stray `GetTensorShapeElementCount` reference fails with > `compute.rs now names 10 ORT API members ([… "GetTensorShapeElementCount" …]), not 9. If you added an FFI call, add an ort_call() beside it and update this table.` A second test pins that the scan finds real members, so it cannot rot into a heuristic that matches nothing and passes forever. **`DispatchAlloc` got the opposite treatment: it is now documented as a lower bound.** The reviewer was right that hand-placed allocation counts cannot be exhaustive — `Vec::new()` does not allocate, `Vec::with_capacity(0)` does not allocate, and whether a `collect` allocates once or twice depends on `size_hint`, not on anything visible at the call site. Claiming exactness there would repeat the same mistake in a place where it cannot be fixed. It is exact for `read_inputs`, where the sites are few and each is pinned; callers needing the true whole-`Run` figure get `CountingAllocator`, a `GlobalAlloc` wrapper to install in a test or bench. This crate deliberately does not install it, since a library defining a `#[global_allocator]` takes that choice away from every dependent. Also fixed: - A node with **no** inputs was charged one allocation for a `Vec::with_capacity(0)` that never happens. Now conditional, pinned by `a_node_with_no_inputs_allocates_nothing`. - `CreateStatus` was counted as a `StatusCreated` but not as an FFI call, undercounting error paths. - **Phases are documented as not being a partition.** On success they are non-exhaustive; on error they *nest*, because a guard closes at scope exit rather than at the early `return`, so `StatusCrossing` opens inside whichever phase was live. `DispatchLookup` is entered twice per node and reports the total. Summing `phase_ns` gives neither wall time nor a pie chart, and a reader who assumed otherwise would be misled. - `probe_is_compiled_out_in_production` no longer implies it has proven `count` is side-effect-free — a disabled build has no storage to observe one. It proves what it can: the guard is a ZST that does not need dropping. Confirmed by the reviewer independently: `mem::zeroed::<OrtApi>()` is sound (all 422 fields are `Option<unsafe extern "C" fn>`, so NPO makes all-zero a valid `None` everywhere); the `without_provenance` fake pointers are never dereferenced; and all four `read_inputs` counts are correct, checked against a counting global allocator. Re-validated after the fixes: clippy clean and `--lib` green in both configurations (253 / 260), and **55/55 `plugin_ort_e2e`** against real ORT. ## What this sets up Two allocations are already visibly waste and are removed in a follow-up that this PR exists to measure: 1. `compute_execute` calls `staging_log(&format!(…))` **unconditionally** — `staging_log` early-returns when disabled, but the `String` is built and three values formatted on every dispatch regardless. 2. `read_inputs` builds `format!("input {i}")` eagerly per input, on the success path, purely to label an error that is usually not raised. Relates to #1077. Independent of #1244/#1246 — additive and feature-gated, so it can land in any order. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Resolves 12 conflicts from the six squash landings of the #1077 stack (#1387, #1409, #1412, #1430, #1433, #1472). compute.rs (6): took main's side throughout. Main is the superset -- #1397's `device_mem_info` now returns `(mem_info, bool)`, #1401 builds one stride entry per *absent* slot rather than cloning every output shape, and the entry/lookup probe spans are new. This branch's `has_absent` guard is subsumed by the per-absent-slot construction. kernel_ctx.rs (6): kept this branch's one-call `GetTensorElementTypeAndShapeDataReference` structure and interleaved main's probe counters into both the borrowed and the legacy path, so `Event::OrtFfiCall` stays a total. Took main's `impl std::fmt::Display` for `validate_dims` (a superset of this branch's `Arguments<'_>`). Kept the inline-rank output-dims path, with the allocation counter moved to the heap-only branch where the allocation actually happens. Two instrumentation updates the merge made necessary: - dispatch_probe's `every_ort_entry_point_is_accounted_for` fired exactly as designed: this branch adds a 12th ORT API member. Table updated to 12 members / 12 `ort_call()` sites. - The pinned costs `1 + 7` calls and `1 + 3` allocations describe the *legacy* path: `fake_api()` zeroes the struct, so the reference hook is null there and the fallback runs. The fast path -- the one real ORT 1.27 takes, and the entire point of this PR -- had no cost pin at all. Added `the_reference_hook_path_costs_exactly_three_ort_calls_per_input` (pins 1 + 3, down from 1 + 7) and `the_reference_hook_path_allocates_less_than_the_legacy_path`. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

justinchuby and others added 4 commits August 19, 2026 04:19

This was referenced Aug 19, 2026

Instrument the routed loop, and fix the O(depth^2) retirement scan (2.31 -> 1.95 at depth 100) #1401

Merged

Reduce cpu ep kernel dispatch overhead #1077

Open

Base automatically changed from squad/resch-dispatch-grid to squad/resch-dispatch-probe August 19, 2026 05:09

justinchuby changed the title ~~Stop asking ORT three times where the inputs live (7 -> 4 FFI calls per Run)~~ Dispatch-overhead stack: benchmark grid, 7 -> 4 FFI calls per Run, routed-loop fixes (2.31 -> 1.71 at depth 100) Aug 19, 2026

justinchuby merged commit f1faf9b into squad/resch-dispatch-probe Aug 19, 2026
6 checks passed

justinchuby deleted the squad/resch-dispatch-meminfo branch August 19, 2026 05:10

justinchuby mentioned this pull request Aug 19, 2026

#1077 dispatch overhead: phase instrumentation, and the fixes it found (2.31 -> 1.71 at depth 100) #1387

Merged

justinchuby mentioned this pull request Aug 19, 2026

perf(plugin): read input type and shape in one ORT call, not five #1246

Merged

Navigation Menu

Repository navigation

Dispatch-overhead stack: benchmark grid, 7 -> 4 FFI calls per Run, routed-loop fixes (2.31 -> 1.71 at depth 100) - #1397

Dispatch-overhead stack: benchmark grid, 7 -> 4 FFI calls per Run, routed-loop fixes (2.31 -> 1.71 at depth 100)#1397
justinchuby merged 5 commits into
squad/resch-dispatch-probefrom
squad/resch-dispatch-meminfo

justinchuby commented Aug 19, 2026 •

edited

Loading

Uh oh!

Uh oh!

github-actions Bot commented Aug 19, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

justinchuby commented Aug 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Stop asking ORT three times where the inputs live

The finding

The change

Why this is behaviour-identical

Evidence

Why counters and not a stopwatch

Also pinned

Validation

Independent review (Opus 4.8, adversarial) — APPROVE WITH NITS, all findings addressed

Claims the reviewer tried to break and could not

Uh oh!

Uh oh!

github-actions Bot commented Aug 19, 2026

🔴 Benchmark Regression Detected

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

justinchuby commented Aug 19, 2026 •

edited

Loading