Skip to content

perf(cpu): run the SiLU, Relu and Clip MLAS routes through run_chunked - #1130

Merged
justinchuby merged 5 commits into
mainfrom
deckard/silu-parallel
Aug 17, 2026
Merged

justinchuby merged 5 commits into
mainfrom
deckard/silu-parallel

Conversation

@justinchuby

@justinchuby justinchuby commented Aug 17, 2026 •

Copy link
Copy Markdown
Owner

Root cause

run_chunked is the seam that #1105 taught to split activation work across the rayon
pool. #1127 fixed dispatch_mlas!, which called its kernel directly and returned,
bypassing that seam entirely — so every MLAS-routed op ran single threaded no matter how
many threads were configured.

The independent review on #1127 reported as a MAJOR finding that three more callers
have the identical pathology, in other files, so they were out of scope there:

caller file
silu_f32_slice kernels/activations.rs
relu_contiguous_f32_mlas kernels/relu.rs:144
Clip (selection) kernels/selection.rs:182
Clip (conv epilogue) kernels/conv.rs:732

mlas-sys documents compute_silu, compute_relu and compute_clip as "Single
threaded; callers shard across threads themselves."
Nobody sharded. This PR wraps all
four call sites in run_chunked.

SiLU needed one extra step. Its MLAS route is followed by a correction scan over the
whole tensor (MLAS's compute_silu is inaccurate outside ±SILU_MLAS_SAFE_BOUND).
Run whole-tensor, that scan streams the buffer a second time from DRAM. It is now
blocked at SILU_CORRECTION_BLOCK = 8192 so each block stays in L2, and the scan is a
branch-free OR-reduction over the input only — the predicate
!x.is_finite() || x.abs() > SILU_MLAS_SAFE_BOUND depends solely on the input, so the
common all-in-band case skips the write loop entirely.

Benchmarks

Session level through the plugin .so, base = origin/main @ b5309f799, 16 threads,
3 interleaved rounds, randomised order, # NOT-ASSIGNED: 0 on every run (no node was
left to ORT's CPU EP). µs, p50.

op n base this PR speedup ORT ORT-rel before ORT-rel after
Clip 1 Mi 351.36 256.98 1.37× 34.73 0.099 0.135
Clip 4 Mi 1284.90 516.92 2.49× 88.91 0.069 0.172
Relu 1 Mi 334.00 252.63 1.32× 40.42 0.121 0.160
Relu 4 Mi 1095.22 494.47 2.22× 78.41 0.072 0.159
Swish 1 Mi 1415.31 404.02 3.50× 226.21 0.160 0.560
Swish 4 Mi 5629.53 549.99 10.24× 474.99 0.084 0.864

Swish (default domain, opset 24) is the ORT-visible spelling of SiLU and is supported
by ORT 1.28, so SiLU does have a real single-node session-level A/B after all — the
earlier note that it did not was wrong, and it is the op that gains the most here.

Kernel-level SiLU, serial_scope vs parallel in-process, 32 threads, so the MLAS route
is compared against itself with only the split changed:

n serial parallel speedup
1 Mi 2051.4 641.6 3.20×
4 Mi 8231.0 1266.0 6.50×
16 Mi 33589.2 3380.0 9.94×

Correctness

  • blocked_correction_matches_the_whole_tensor_loop_bit_for_bit — the blocked,
    OR-reduced scan is compared bit for bit against the original whole-tensor loop over
    in-band values, out-of-band values, ±Inf, NaN, ±0, denormals and values sitting
    exactly on SILU_MLAS_SAFE_BOUND, at lengths that straddle the block boundary.
  • silu_reaches_run_chunked_parallel_branch — asserts the mechanism, not the
    output, using the PARALLEL_DISPATCHES counter added in perf(cpu): run the MLAS activation routes through run_chunked #1127. Verified to falsify:
    reverting the run_chunked wrapper makes it fail.
  • silu_is_thread_count_invariant — identical results across pool sizes.
  • No tolerance was relaxed anywhere. No numerical behaviour changes: this PR only
    changes who runs the arithmetic, plus a blocking/reduction rewrite that is proven
    bit identical.

cargo test -p onnx-runtime-ep-cpu --features mlas --lib → 1322 passed, 0 failed.
Both feature configurations build. cargo fmt clean.

Limitations

  • Clip, Relu and SiLU still lose to ORT at these sizes (0.135–0.864×). This PR is a
    2.2–10.2× step toward the architectural requirement that our CPU EP beat ORT on every
    op it accepts; it does not finish the job, and no fallback was added. The remaining
    gap is the general 16-thread scaling gap tracked in docs/performance/CPU_ACTIVATION_GAPS.md
    — ORT scales these ops ~14× from 1→16 threads, we manage ~6×, because we split over
    our own rayon pool rather than ORT's intra-op pool. The host_parallel seam over
    KernelContext_ParallelFor is the next step.
  • 1 Mi is exactly PAR_MIN_LEN, so gains there are smaller and noisier than at 4 Mi.
  • 16-thread medians on this shared machine are noisy; untouched control ops swung up to
    36% across 3 rounds. The 4 Mi wins are far outside that band. The 1 Mi Clip/Relu
    numbers are closer to it and should be read as directional.
  • run_chunked, PAR_MIN_LEN and parallel_dispatches are widened to pub(crate)
    because the three other callers live in sibling modules.

deckard and others added 3 commits August 17, 2026 15:02
The MLAS routes for SiLU, Relu and Clip called their kernel directly
instead of going through `run_chunked`, so they never split across the
rayon pool no matter how many threads were configured. This is the same
pathology #1127 fixed for the `dispatch_mlas!` ops, reported as a MAJOR
finding by the independent review on that PR.

`mlas-sys` documents `compute_silu`, `compute_relu` and `compute_clip`
as single threaded, with sharding left to the caller. Nobody sharded.

SiLU additionally needed its correction scan blocked so the parallel
chunks stay in L2, and the scan is now a branch-free OR-reduction over
the input alone, so the common all-in-band case skips the write loop.

Measured at 16 threads, session level through the plugin, zero nodes
left unassigned:

  Clip  4 Mi  1284.90 ->  516.92 us  2.49x
  Relu  4 Mi  1095.22 ->  494.47 us  2.22x
  Swish 4 Mi  5629.53 ->  549.99 us  10.24x

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Independent review findings on #1130.

MAJOR: `.work_backup.rs`, a 988 line mid-development snapshot of
activations.rs, had been committed at the repo root by accident. Removed.

MINOR: the `silu_bench` doc comment claimed Silu/Swish have no single-op
equivalent in ORT, which this PR itself disproves — ORT 1.28 implements
`Swish` in the default domain at opset 24, and the PR reports a
session-level A/B built on it. Comment corrected.

NIT: `silu_reaches_run_chunked_parallel_branch` hard-failed on a
single-core runner instead of skipping. It now returns early with a note.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The re-review noted that `assert_same` and `assert_parallelises` in
simd_activations.rs carry the same unguarded hard assert that was just
fixed for `silu_reaches_run_chunked_parallel_branch`. They predate this
PR, but they are its closest analogs, so they get the same early-return
skip rather than a spurious failure on a single-core runner.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@codecov

codecov Bot commented Aug 17, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 63.63636% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.60%. Comparing base (2b4eaeb) to head (6b092e5).
⚠️ Report is 2 commits behind head on main.

Files with missing lines Patch % Lines
...nnx-runtime-ep-cpu/src/kernels/simd_activations.rs 60.00% 4 Missing ⚠️
Additional details and impacted files

Impacted file tree graph

@@            Coverage Diff             @@
##             main    #1130      +/-   ##
==========================================
+ Coverage   79.94%   80.60%   +0.66%     
==========================================
  Files         368      366       -2     
  Lines      160607   157585    -3022     
  Branches   160607   157585    -3022     
==========================================
- Hits       128396   127028    -1368     
+ Misses      27494    25853    -1641     
+ Partials     4717     4704      -13     
Flag Coverage Δ
cli-ort-linux 83.79% <ø> (ø)
cli-ort-windows ?
mlas ?
offline 80.49% <63.63%> (+0.80%) ⬆️

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

Files with missing lines Coverage Δ
...tes/onnx-runtime-ep-cpu/src/kernels/activations.rs 91.92% <100.00%> (ø)
crates/onnx-runtime-ep-cpu/src/kernels/relu.rs 90.74% <ø> (ø)
...rates/onnx-runtime-ep-cpu/src/kernels/selection.rs 82.64% <ø> (ø)
...nnx-runtime-ep-cpu/src/kernels/simd_activations.rs 97.23% <60.00%> (-0.23%) ⬇️

... and 11 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@justinchuby
justinchuby marked this pull request as ready for review August 17, 2026 16:34
@justinchuby
justinchuby merged commit 34095af into main Aug 17, 2026
10 of 17 checks passed
@justinchuby
justinchuby deleted the deckard/silu-parallel branch August 17, 2026 16:34
@github-actions

Copy link
Copy Markdown

🔴 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
🔴 add/medium_f16_threads=1-internal/262144 117.69 µs 163.60 µs +39.0%
⚠️ qwen3_sampling_processors/top_k_top_p_full_sort_baseline 6.40 ms 8.28 ms +29.2%
⚠️ sampling_latency/greedy_per_token 2.99 µs 3.75 µs +25.1%
⚠️ tokenization/decode_tokens_per_second 5.88 ms 7.12 ms +21.1%
⚠️ sampling_latency/top_k_per_token 48.26 µs 56.71 µs +17.5%
⚠️ matmul/medium_generic_bf16_threads=8/32x512x512 505.41 µs 587.47 µs +16.2%
✅ add/medium_bf16_threads=1-internal/262144 106.20 µs 121.19 µs +14.1%
✅ tokenization/encode_tokens_per_second 347.03 µs 385.32 µs +11.0%
✅ gather/medium_bf16_threads=1-internal/32768 2.19 µs 2.40 µs +9.5%
✅ add/small_bf16_threads=1-internal/1024 461.5 ns 504.8 ns +9.4%
✅ gather/small_f16_threads=1-internal/4096 447.5 ns 484.9 ns +8.4%
✅ qwen3_sampling_processors/top_p_fast_after_top_k 560.67 µs 601.58 µs +7.3%
✅ gather/small_bf16_threads=1-internal/4096 471.4 ns 499.8 ns +6.0%
✅ gather/small_f32_threads=1-internal/4096 649.7 ns 684.6 ns +5.4%
✅ gather/medium_f32_threads=1-internal/32768 3.51 µs 3.64 µs +3.6%
✅ sampling_latency/top_p_per_token 386.83 µs 399.80 µs +3.4%
✅ add/large_bf16_threads=1-internal/4194304 1.72 ms 1.77 ms +2.8%
✅ grammar_masking/llguidance_compute_mask/32 75.79 µs 77.82 µs +2.7%
✅ qwen3_sampling_processors/top_k_top_p_fast 728.49 µs 740.18 µs +1.6%
✅ gather/medium_f16_threads=1-internal/32768 2.21 µs 2.22 µs +0.5%
✅ reduce_mean/large_f32_threads=1-internal/262144 1.13 ms 1.14 ms +0.2%
✅ matmul/small_generic_f32_threads=1/1x256x256 34.00 µs 34.04 µs +0.1%
✅ add/large_f16_threads=1-internal/4194304 1.83 ms 1.81 ms -0.8%
✅ add/small_f16_threads=1-internal/1024 456.8 ns 452.2 ns -1.0%
✅ matmul/small_generic_f32_threads=8/1x256x256 32.70 µs 32.31 µs -1.2%
✅ block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1 456.36 µs 450.90 µs -1.2%
✅ reduce_mean/medium_f32_threads=1-internal/65536 243.11 µs 240.15 µs -1.2%
✅ matmul/large_generic_f32_threads=8/32x1024x1024 4.28 ms 4.19 ms -2.1%
✅ add/large_f32_threads=1-internal/4194304 785.82 µs 760.63 µs -3.2%
✅ matmul/small_generic_bf16_threads=8/1x256x256 31.89 µs 30.65 µs -3.9%
✅ kv_cache/alloc_dealloc_pages 40.72 µs 38.91 µs -4.5%
✅ logit_processing/seven_processor_chain_per_step 341.65 µs 325.66 µs -4.7%
✅ block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1 163.23 µs 155.00 µs -5.0%
✅ sampling_latency/min_p_per_token 226.72 µs 213.75 µs -5.7%
✅ reduce_mean/small_f32_threads=1-internal/4096 15.41 µs 14.50 µs -5.9%
✅ matmul/small_generic_f16_threads=1/1x256x256 29.77 µs 27.84 µs -6.5%
✅ gather/large_f32_threads=1-internal/131072 25.49 µs 23.83 µs -6.5%
✅ qwen3_sampling_processors/top_k_full_sort_baseline 2.40 ms 2.23 ms -7.1%
✅ matmul/medium_generic_bf16_threads=1/32x512x512 565.12 µs 522.99 µs -7.5%
✅ qwen3_sampling_processors/top_k_partial_selection 168.79 µs 154.59 µs -8.4%
✅ add/medium_f32_threads=1-internal/262144 28.97 µs 26.28 µs -9.3%
✅ matmul/large_generic_f32_threads=1/32x1024x1024 10.05 ms 9.00 ms -10.4%
✅ add/small_f32_threads=1-internal/1024 239.1 ns 206.5 ns -13.6%
🟢 block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024 612.61 µs 516.98 µs -15.6%
🟢 matmul/small_generic_bf16_threads=1/1x256x256 36.90 µs 30.65 µs -16.9%
🟢 qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline 4.66 ms 3.82 ms -18.0%
🟢 matmul/small_generic_f16_threads=8/1x256x256 39.29 µs 31.82 µs -19.0%
🟢 gather/large_f16_threads=1-internal/131072 11.63 µs 9.20 µs -20.9%
🟢 matmul/medium_generic_f16_threads=8/32x512x512 38.03 µs 29.71 µs -21.9%
🟢 matmul/large_generic_f16_threads=1/32x1024x1024 94.02 µs 73.08 µs -22.3%
🟢 gather/large_bf16_threads=1-internal/131072 13.23 µs 10.15 µs -23.3%
🟢 matmul/medium_generic_f32_threads=1/32x512x512 3.01 ms 2.23 ms -25.7%
🟢 matmul/medium_generic_f16_threads=1/32x512x512 43.95 µs 32.44 µs -26.2%
🟢 matmul/large_generic_bf16_threads=8/32x1024x1024 1.76 ms 1.28 ms -27.3%
🟢 matmul/large_generic_bf16_threads=1/32x1024x1024 2.80 ms 1.81 ms -35.2%
🟢 matmul/large_generic_f16_threads=8/32x1024x1024 124.84 µs 80.50 µs -35.5%
🟢 matmul/medium_generic_f32_threads=8/32x512x512 1.77 ms 906.30 µs -48.9%
🟢 block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024 78.92 µs 40.20 µs -49.1%
🟢 block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024 79.75 µs 37.50 µs -53.0%

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

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

justinchuby added a commit that referenced this pull request Aug 17, 2026
## What happened

#1130 (mine) wrapped the `Clip` MLAS call in `run_chunked` so it would
use the
thread pool. `run_chunked` is generic, so doing that from `selection.rs`
created a new
instantiation in a module that had never had one.

The runtime path of every other unary op was untouched — not one
instruction — and all
1332 tests still passed. But the crate's codegen units repartitioned and
the AVX2 unary
kernels in `simd_activations.rs` stopped being vectorised.

This shipped. It was found while collecting numbers for the final
report: `Sqrt` had
gone from beating ORT by 1.9× to losing at 0.7×, and the ORT-relative
position of five
ops had collapsed in a way no code change explained.

## Evidence

Bisected by rebuilding `main` with one file at a time reverted to
`1f1ce4b74`
(the commit before #1130). n = 65536, 1 thread, `taskset -c 8-23`, 4
interleaved
rounds, µs p50. **Only `selection.rs` matters** — reverting `relu.rs` or
`conv.rs`
alone restored nothing:

| op | main (#1130) | revert relu.rs | revert **selection.rs** | revert
conv.rs | pre-#1130 |
|---|---:|---:|---:|---:|---:|
| Sqrt | 48.5 | 48.4 | **21.3** | 44.6 | 21.5 |
| Tanh | 55.2 | 57.0 | **31.6** | 60.0 | 31.8 |
| Sigmoid | 56.9 | 55.6 | **31.0** | 57.1 | 31.1 |
| QuickGelu | 64.5 | 64.5 | **42.5** | 64.2 | 42.5 |
| FastGelu | 79.6 | 79.3 | **55.5** | 79.3 | 55.8 |
| Erf | 62.2 | 62.4 | 62.2 | 62.2 | 62.5 |
| Relu | 21.3 | 26.5 | 22.6 | 23.0 | 20.4 |

Reproduced independently by building the same commit in a second
worktree, so it is not
a build-directory artefact. Adding `#[inline]` to `run_chunked` did
**not** help, which
rules out a plain inlining decision and points at codegen-unit
partitioning.

## The diagnosis, confirmed mechanically

If the cause really is codegen-unit partitioning, then forcing the crate
into a single
codegen unit should erase the regression with no source change at all.
It does.
`CARGO_PROFILE_RELEASE_CODEGEN_UNITS=1` on the *unfixed* commit
`34095af0f`, same
machine, 4 interleaved rounds, µs p50:

| op | n | main, default CGUs | main, `codegen-units=1` | this PR,
default CGUs |
|---|---:|---:|---:|---:|
| Sqrt | 64 Ki | 48.3 | **21.6** | 21.2 |
| Sqrt | 1 Mi | 667.4 | **261.4** | 274.4 |
| Sigmoid | 1 Mi | 813.3 | **504.8** | 441.6 |
| QuickGelu | 1 Mi | 929.4 | **575.8** | 573.0 |
| FastGelu | 1 Mi | 1171.2 | **788.2** | 808.5 |

So the diagnosis is not inferred from a bisect alone — the proposed
mechanism, applied
directly, reproduces the cure.

`codegen-units = 1` (or LTO) in the release profile would remove this
whole fragility
class permanently, and is the more durable answer. It is deliberately
**not** in this PR:
it is a workspace-wide build-policy change that affects every crate and
every
contributor's build time (the plugin alone went 14 s to 63 s here), it
would need its own
measurement across the whole EP rather than the activation kernels, and
it does not
belong in a regression fix. Filed as the follow-up this PR's limitation
section points
at. The source fix costs nothing and is independent of it.

## The fix

`run_chunked` is private to `simd_activations.rs` again — the compiler
now enforces the
rule, not a convention. Callers elsewhere go through one of two entry
points that are
instantiated *in that module*:

- `run_chunked_fn(input, output, body: fn(&[f32], &mut [f32]))` —
deliberately a `fn`
pointer, not `impl Fn`, so every caller shares one instantiation. Used
by `Relu` and
  `SiLU`.
- `clip_chunked(input, output, min, max)` — `Clip` needs captured
bounds. It takes the
serial decision itself so the short case is a direct call rather than
one through a
closure the optimiser can no longer see into; without that, `Clip`
itself paid 12%.

## Result

n = 65536 and 1 Mi, 1 thread, 5 interleaved rounds, µs p50:

| op | n | main (#1130) | this PR | pre-#1130 |
|---|---:|---:|---:|---:|
| Sqrt | 64 Ki | 48.5 | **21.4** | 21.7 |
| Sqrt | 1 Mi | 667.4 | **259.2** | 275.1 |
| Tanh | 1 Mi | 776.1 | **437.8** | 436.4 |
| Sigmoid | 1 Mi | 807.6 | **440.1** | 502.2 |
| QuickGelu | 1 Mi | 923.1 | **574.7** | 607.0 |
| FastGelu | 1 Mi | 1166.9 | **785.3** | 788.1 |
| Clip | 1 Mi | 268.7 | **267.5** | 267.4 |
| Swish | 1 Mi | 728.8 | **720.8** | 1385.6 |

Everything is back to its pre-#1130 level **and** #1130's own win is
kept: `Swish` is
still 1.92× faster than before #1130, and `Clip`/`Relu` still reach the
pool at
≥ `PAR_MIN_LEN`.

## Regression guard


`chunking_instantiation_is_local::no_module_outside_this_one_instantiates_run_chunked`
walks the crate source and fails if any module other than
`simd_activations.rs`
instantiates `run_chunked`, naming the offending file and line.

Verified to falsify: re-widening the visibility and pointing
`relu.rs:144` back at
`run_chunked` makes it fail with
`Offending call sites: ["…/kernels/relu.rs:144"]`.

This class of bug produces no wrong answers, no test failures and no
diff in the file
that slows down, so a mechanical guard is the only thing that catches
it.

## Tests

- `cargo test -p onnx-runtime-ep-cpu --features mlas --lib` → **1332
passed, 0 failed**
- `cargo test -p onnx-runtime-ep-cpu --lib` → **1305 passed, 0 failed**
- `cargo fmt --all` clean.
- No numerical change: `clip_chunked`'s serial branch calls exactly the
function the
  closure called, and `run_chunked_fn` forwards unchanged.

## Limitations

- The mechanism is codegen-unit partitioning, which the compiler makes
no promises
about. The guard encodes the rule that was measured to work on this
toolchain; it
cannot prove the next refactor is safe. That is why the guard names the
symptom and
  the measured cost in its failure message.
- Measured on one machine (AMD EPYC 9V74, AVX2/FMA/F16C,
`taskset`-pinned, 5 rounds).
The *direction* is unambiguous — up to 2.3× on `Sqrt` and roughly
1.5–1.8× on the
  other four — but the exact figures are not portable.

---------

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants