Repository navigation
perf(cpu): stop building the decode pool for decode-shaped borrowed int4 - #1434
Conversation
The bounded, pinned decode Rayon pool is built eagerly on the first `MatMulNBits` projection and then, for a pure-decode workload, never runs any model arithmetic: `ONNX_GENAI_PROFILE_OPS` attributes 99.95% of the forward pass to a single `MatMulNBits` per iteration, and the task runtime takes 1.01 dispatches per iteration for it. The pool is a pass-through -- 16 resident workers whose only cost is the per-dispatch install/wake. That cost is not idle spin. Holding the iteration count fixed and stretching the inter-token gap 100us -> 4000us grows wall time 2.9x while the pool's CPU stays flat at ~350ms, so it is ~1.8ms of CPU per dispatch paid to enter and leave a pool that computes nothing. Defer building it. `with_decode_pool_lazy` runs the kernel inline while publishing the width the pool *would* have had, and `parallel_output_rows_repeated` installs the pool on demand if a fan-out actually routes `Wide`. Routing is therefore unchanged: `wide` reads the same value it read inside the installation, so no shape moves between the task runtime and Rayon. Scope is what makes this provable rather than hopeful. The deferral is gated on `m == 1`, and at `m == 1` every kernel reachable from the borrowed int4 closure parallelises solely through `parallel_output_rows_repeated`. The one branch that drives `par_chunks_mut` directly, `borrowed_affine_int4_matmul_prefill`, is gated on `m >= 2` and so is unreachable; prefill keeps installing eagerly, as do the seven other call sites whose kernels reach Rayon by other routes (`packed_nbits_gemm`, `int8_matmul`, `parallel_n16_output_rows`, `parallel_kai_output_rows`). A one-shot probe at every `with_decode_pool` call site confirms this one is the only builder for a decode workload. Measured at a 16-core budget, interleaved against the unmodified binary: the process drops from 48 threads to 32, the 16 `onnx-genai-decode-*` workers and their 340-490ms disappear entirely, and voluntary context switches per iteration fall from ~18.7 to ~14.8. Latency is not claimed here: the host was at load average 26 for these runs, well outside the noise guard, so p50/p90 are deferred to a quiet window. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ust the route Opus review caught that `output_chunk_len` reads `rayon::current_num_threads()` directly, so it was the one input to the fan-out that the deferral hint did not cover. The executor choice was preserved but the *grain* was not: with no explicit budget the global pool is `available_parallelism()` wide while the decode pool is narrower, so a deferred fan-out would partition for the wrong pool and, at the boundary, flip serial and parallel. Results stay correct either way -- row sharding is associative -- but "routing is unchanged" was not literally true. Route both the executor choice and the grain through one `effective_fan_out_width()`, so a deferred fan-out reproduces the installed case exactly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Opus review: no bugs, one nuance closedIndependent Opus review traced every kernel reachable from the three converted sites at It did find one real imprecision, now fixed in
It rated this non-blocking (unreachable in the production scoped path, and never changes results because row sharding is associative), but it made the PR's own "routing is unchanged" claim not literally true, so I closed it rather than caveat it. Both the executor choice and the grain now go through one Latency: measured on a quiet window, and it is neutralThe Draft deferred p50/p90 because the host was at load average 26. It quieted (cpu/wall back to 9.5–10.0, in band), so here is the interleaved A/B with alternating arm order, 400 iters, 16-core budget, plus an A/A null control run in the same window: A/A null control (identical binaries, separate processes): p50 ratios 1.000, 0.828, 0.997 → null band 0.83–1.21. A/B (gated / baseline), alternating order:
Median 1.073, every rep inside the null band. So: no detectable latency change in either direction. I am not claiming a speedup, and there is no regression. The win is structural, and those metrics are load-insensitive. From a clean census pair in the same window (cpu/wall 9.56 vs 9.53):
Why the 350 ms does not reappear elsewhereWorth stating explicitly, because the thread census looks like it moved: the caller group's CPU rises by roughly what the decode pool loses. That is attribution, not new work. Taking out of Draft: Opus green, full |
## What Unbreak the two CI lanes that are **red on `main` right now**. Three independent breakages, all pre-existing and all reproduced on an unmodified `dbade34c1` checkout with the same stable 1.97.1 toolchain CI installs: | # | gate | breakage | fix | |---|------|----------|-----| | 1 | `cargo fmt --all -- --check` | 5 sites / 4 files | rustfmt | | 2 | `Rust quality` clippy | `clippy::unnecessary_map_or` — `dispatch.rs:26` | `map_or(true, f)` → `is_none_or(f)` | | 3 | `Fast (Linux x86_64)` clippy `--all-targets` | `clippy::inconsistent_digit_grouping` — `cost-model/model.rs:314` | `2_000_000_000_000_0` → `20_000_000_000_000` | Both lanes build with `RUSTFLAGS: -D warnings`, so #2 and #3 are hard errors, not warnings. **Every PR that merges `main` inherits all three** — verified on #1434 and #1420. Nothing in the queue can go green until this lands. #3 is worth calling out: it is invisible to a plain `cargo clippy` because the literal lives in a `#[cfg(test)]` module. Only the `--all-targets` invocation in the Fast lane sees it. ## Semantics Both non-fmt changes are provably value-preserving: - `is_none_or(f)` is the rewrite the lint itself suggests, and is definitionally `map_or(true, f)`: `None` → `true`, `Some(v)` → `f(v)`. `gqa_shape_capacity_bound_enabled()` is unchanged — unset stays enabled, the falsey spellings stay disabled. - `20_000_000_000_000 == 2_000_000_000_000_0` (both 2e13), which is what the test's own comment already claims — *"2e13 FLOP / 2e13 = 1 s"*. `op_cost_takes_roofline_max` still asserts the compute term dominates. ## Validation Ran locally per the delayed-Actions directive, on this head merged with `origin/main` @ `dbade34c1`: | gate | result | |------|--------| | `cargo fmt --all -- --check` | **0 diffs** | | `cargo clippy --locked --all-targets $(workspace_test_packages.py cargo-args offline-linux) -- -D warnings` | **exit 0** | | `cargo test --locked $(… offline-linux)` | **3943 passed, 0 failed, exit 0** | | `scripts/check_cross_compile.sh` | **PASS** — x86_64 + aarch64 full offline set | | `benchmark_muse_native_local.py --self-test --require-numpy` | 43 cases passed | | `check_publish_order.py` / `check_profile_table.py` / `check_platform_naming.py` | PASS | | `check_dispatch_reachability.py` / `check_feature_gate_coverage.py` | PASS | | `check_dispatch_manifest.py` (`--self-test` and plain) | PASS | | `workspace_test_packages.py verify` | PASS | | `verify_documented_env_vars.py` | PASS — 113 documented, 13 known-unimplemented | | MLAS cfg: `-p onnx-runtime-ep-cpu --no-default-features --features mlas` | `moe::` 19 passed · `qlinear_matmul::` 30 passed · `optimization_registry_excludes_nchwc_without_cnn_ops` 1 passed | | `cargo clippy -p onnx-genai-engine --features native-backend` | clean | | `cargo build -p onnx-runtime-ep-cpu-plugin --features mlas` | clean | ### Windows ARM64: not validated locally — stated as a blocker, then bounded I could not run `Rust (Windows ARM64)` here and I am **not** claiming it as a pass. Two routes were attempted, both fail *identically on unmodified `main`*, so neither can discriminate this PR from baseline: - **`cargo-xwin` / clang-cl** — installed, MSVC CRT + SDK downloaded, correctly targeting `aarch64-pc-windows-msvc`. Fails in vendored `mlasi.h` on NEON intrinsics (`veorq_s32`, `vdupq_n_f32`, …) that MSVC supplies but clang-cl in MSVC mode does not. - **`aarch64-unknown-linux-gnu` + GNU cross toolchain** as an ARM64-NEON proxy — gets much further, compiles most of the ARM64 MLAS source set, then fails on `activate_fp16.cpp`. What makes this safe to merge anyway is **dependency-graph disjointness**, not a judgement call. That lane builds only `mlas-sys` and `onnx-runtime-ep-cpu-plugin --features mlas`. This PR touches `onnx-genai-engine` (tests), `onnx-runtime-ep-cuda` and `onnx-runtime-session`: ``` $ cargo tree -p onnx-runtime-ep-cpu-plugin --features mlas -e normal --prefix none \ | sort -u | grep -cE "onnx-runtime-ep-cuda|onnx-runtime-session|onnx-genai-engine" 0 ``` Zero of the crates this PR modifies are in that lane's graph, so it cannot observe this change. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Local validation (delayed-Actions directive)Head merged with The 48 -> 32 reduction, re-measured from scratch on latest mainSame worktree, same release binary recipe, same model, same host. The only Model
The So the decode pool is removed by construction on this path, not merely That the 4 and 8 rows are unchanged is the expected shape, not a miss: below Scope honestyThe saving is specific to the covered path. On the same model with the decode Behavioural falsifiers
Repository gates
|
🔴 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
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #1434 +/- ##
==========================================
+ Coverage 82.10% 82.64% +0.54%
==========================================
Files 12 12
Lines 5471 5475 +4
Branches 5471 5475 +4
==========================================
+ Hits 4492 4525 +33
+ Misses 780 757 -23
+ Partials 199 193 -6
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
1 participant
Building the decode pool is unconditional today: every borrowed-int4 call installs a 16-worker Rayon pool before dispatching, including the decode-shaped calls that then route all of their work to the task runtime and never touch it. Those workers are constructed, parked, and torn down once per call for nothing.
This defers the construction for exactly one path — decode-shaped (
m == 1) borrowed int4 — and hands the routing logic the width it would have installed, so the executor choice and the partition grain are both unchanged.What is actually covered
The scope is deliberately narrow, and the narrowness is what makes it provable rather than measured:
packed_nbits_gemv,gemv_nk, and the borrowed-int4 site gated onm == 1.parallel_output_rows_repeated(parallel_output_rowsdelegates to it), which is the hookable dispatcher.m == 1gate is what makes this safe by construction, not by inspection:borrowed_affine_int4_matmul_prefilldrives Rayon directly, and it is unreachable atm == 1. The other sevenwith_decode_poolsites are untouched, so the kernels that use Rayon directly —packed_nbits_gemm,int8_matmul,int8_row,parallel_n16_output_rows,parallel_kai_output_rows— keep eager installation.Routing is preserved rather than assumed.
flat_fan_outbranches onrayon::current_num_threads(), which equalled the decode width only because we were installed;effective_fan_out_width()reproduces exactly that width while deferred. The second commit extends the same helper tooutput_chunk_len, so the grain cannot drift from the installed grain either — without it a deferred call partitioned into 4096-row chunks where the installed one used 64.Five tests, two of them verified as falsifiers by deliberately breaking the thing they check:
m >= 1makesonly_the_decode_shaped_borrowed_int4_path_defers_the_poolfail;a_deferred_fan_out_partitions_for_the_pool_it_would_have_installedfail (chunk 64 vs 4096).Measured
16-core budget, interleaved with alternating arm order in a single session.
onnx-genai-decode-*workersLatency is neutral, and that is the claim — not an improvement. Six A/B reps gave 0.938, 1.180, 1.130, 1.115, 1.031, 1.026 (median 1.073), and an A/A null control measured in the same window gave a 0.83-1.21 band. Every ratio is inside the band, so this PR does not demonstrate a latency change in either direction. The win is structural: 16 fewer threads and a third fewer voluntary context switches.
One number deserves an explicit caveat: total CPU is flat, not lower. The decode pool's CPU does not disappear, it reappears on the caller thread. That is attribution changing, not work being removed — the work was always the caller's, it was just being done by borrowed workers.
onnx-runtime-ep-cpulib suite: 1453 passed, 0 failed on merged latest main. fmt and clippy clean for this crate. (cargo fmt --checkcurrently reports five sites repo-wide, all inherited from main and none in the one file this PR touches; #1393 repairs them.)