perf(cpu-ep): vectorise the production softmax exp on AVX2+FMA by justinchuby · Pull Request #1234 · justinchuby/onnx-genai

justinchuby · 2026-08-18T07:49:44Z

Why

The CPU EP that ships does not link MLAS. Its softmax evaluated one libm f32::exp per element, and #1231 - which finally let the benchmark build the arm that actually ships - showed what that cost: 9.3-10.9x ORT across the entire softmax family, the single largest gap the CPU EP had. scale_mask_softmax_serial routes through the same helper, so the fused attention path paid it too.

The doc comment on softmax_rows_in_place had claimed an "8-lane polynomial". That was only ever true of the MLAS arm; #1231 corrected the comment, and this PR makes it true of the arm we ship.

What

Both the in-place and out-of-place non-MLAS forms now route every row through one softmax_row_core(src, dst, d). On AVX2+FMA it evaluates exp eight lanes at a time:

Cody-Waite split of ln2 so x - k*ln2 keeps its low bits;
degree-6 Cephes minimax polynomial, Horner form, all FMAs;
2^k built directly as (k + 127) << 23.

Dispatch is a runtime is_x86_feature_detected!("avx2") && ("fma"); the original scalar loop stays for non-AVX2 targets and for sub-8 tails. All loads and stores are loadu/storeu, so no alignment is assumed. Every 8-lane block is loaded from src before the matching block is stored to dst, so an exact in-place alias is sound - and output_direct_write_eligible already refuses the direct-write slice on partial overlap, so nothing else can reach it.

Sharing one core between the two forms is what keeps them bit-identical, which out_of_place_softmax_is_bit_identical_on_pathological_rows (from #1219) pins.

Non-finite lanes

Handled by explicit masks rather than by trusting the polynomial:

NaN: unordered compare plus blendv, so a poisoned row still yields a NaN sum and a NaN row.
-inf and deep underflow: ordered <= -87.336 compare plus andnot, forcing exactly 0.

The clamp max(x, -87.336) also keeps -inf/NaN out of the exponent build, so cvtps_epi32 only ever sees k in [-126, 0] and the biased exponent stays valid.

Flushing at -87.336 cannot produce a 0/0 row: x = v - row_max, so the max lane is always exp(0) = 1 and the sum is at least 1. The only way every lane underflows is a fully masked row, whose max is -inf - and that row goes down the NaN path instead, exactly as §29 documented.

Both masks are falsifiable. Disabling the NaN blend fails out_of_place_softmax_is_bit_identical_on_pathological_rows; disabling the underflow andnot fails large_and_masked_logits_stay_finite. I ran both.

Measured

Production-native arm only (--features bench-native, nm -C reports 0 MLAS symbols in the binary). Arms interleaved innermost, 8 reps, min-of-ratio. Load average on the box was ~10 from other work, which inflates both arms - the ratio between arms is the claim here, not the absolute column.

model	before native/ort	after native/ort	speedup	parity
sm_prefill_h32_s512	9.044	2.005	4.51x	PASS
sm_bert_b8_s128	12.909	3.411	3.78x	PASS
sm_decode_h32_kv1024	11.067	1.412	7.84x	PASS
sm_decode_h32_kv2048	18.333	2.556	7.17x	PASS
sm_decode_h32_kv4096	31.286	4.286	7.30x	PASS
sm_decode_h32_kv8192	48.926	6.962	7.03x	PASS

The least-contended cell (kv1024, 1.412) matches an earlier quiet-machine measurement of this change at 1.39-1.52, so I read the true post-fix position as ~1.4-1.5x ORT, not parity. That is a 6-7x improvement and it is still a loss; §29's remaining three-pass memory-bound structure is the next lever. I am not claiming softmax is won.

Miri

Worth recording precisely, because it is easy to assume the opposite.

Miri reports is_x86_feature_detected!("avx2") == false, so Miri never executes this kernel - it only ever runs the scalar fallback. It cannot validate the new intrinsics, and no amount of adding softmax to the Miri job would change that.
Miri deliberately returns a nondeterministic approximation for libm calls. Two evaluations of (-0.5f32).exp() in one Miri process gave 0x3f1b4595 and 0x3f1b4599. It also picks a nondeterministic NaN sign/payload per operation - under -Zmiri-seed=1 the masked row came back 0xffc00000 from one arm and 0x7fc00000 from the other, and seed 2 swapped which arm produced which.

So any bit-identity assertion over exp fails under Miri by construction, on a seed-dependent lane. The test is now #[cfg_attr(miri, ignore)] with that reasoning recorded inline. Hardware exp and hardware NaN propagation are both deterministic, which is why it holds off Miri. CI's Miri job scopes ep-cpu to strided/provider/dtype, so this was never red in CI.

Separately, and not touched here: Miri finds genuine pre-existing UB at dtype.rs:838, an unaligned &[f32] built by to_dense_f32_widen, reachable from kernels::softmax::tests::softmax_axis0 and kernels::attention::tests::qk_matmul_output_mode3_is_softmax. It reproduces on main without this PR. I will fix it separately.

Validation

cargo test -p onnx-runtime-ep-cpu --lib - 1341 passed, 0 failed.
cargo test -p onnx-runtime-ep-cpu --lib --features mlas softmax - 29 passed.
cargo test -p onnx-runtime-ep-cpu --lib --no-default-features --features mlas kernels::moe:: (a real CI lane) - 18 passed.
cargo clippy -p onnx-runtime-ep-cpu --lib --all-targets and the same with --no-default-features --features mlas - both clean.
cargo fmt --all --check - clean.
Parity PASS in all 6 models x 8 reps x both arms.

Review

Opus review returned MERGE WITH FIXES and found a blocking regression I had introduced. Both fixes are in.

Blocking, fixed. My clippy fix for fix(bench): let bench_generic build without MLAS, and label the arm it measured #1231 moved the mlas feature guard into a const block. Cargo features are additive and indistinguishable at cfg time, so a compile-time assert cannot tell "mlas became a default" from "mlas was explicitly requested" - it fires on both. CI lanes that build this crate with --features mlas (the MoE and QLinearMatMul reference-arm tests, and the MLAS-only feature-configuration check) would have failed to compile the test target with E0080 instead of running. Restored the runtime assert with an #[allow], and recorded in a comment why a const block is wrong here rather than merely lint-noisy.

The reviewer also caught the process error behind it: I ran the --features mlas validation before making that edit and never re-ran it, so my own PR body carried a validation claim that did not hold. That claim is corrected above. This is the same class of mistake as the withdrawn measurements in §30 - validating a state that is not the state being shipped.
Coverage gap, fixed. The bit-identity test only used D = 8, a whole vector body with no scalar tail. It now runs at 8, 11 and 6: a full body, a body plus a 3-lane tail, and a pure tail with no vector body. This matters because a NaN in a vectorized lane is canonicalized by the mask while a NaN in a tail lane is whatever libm returns, so a change that stopped the two forms sharing one core would diverge on exactly these rows. The falsifier fires at all three widths.

The reviewer independently confirmed, with numbers, several things I had asserted: max 1 ULP error for the new exp against an f64 reference over the whole [-87.336, 0] domain (and over [0, 40], outside the softmax domain); worst normalized-softmax relative error 1.6e-6 vs f64; that the flush threshold lands exactly on the normal/subnormal boundary, discarding at most 1.18e-38 against a row sum of at least 1; that the row max is taken over body and tail; that output_direct_write_eligible rejects partial overlap so the out-of-place path always gets disjoint buffers; and that cvtps_epi32 cannot disagree with the earlier round_ps because k is already integer-valued. They also reproduced both falsification experiments, and found the NaN-mask one is stronger than I claimed - it additionally fails a_fully_masked_row_reproduces_the_reference_runtimes_nan, because the max(x, underflow) clamp turns NaN into -87.336 and the row then computes ~1e-38 instead of NaN.

No ORT fallback, no defer, no runtime MLAS.

The non-MLAS softmax called f32::exp() once per element, one libm call per score, which was the whole 9-11x gap against ORT's CPU softmax on the attention shapes. Route both the in-place and out-of-place non-MLAS paths through one shared row reducer (softmax_row_core) that evaluates exp eight lanes at a time on x86 AVX2+FMA (Cephes expf: k*ln2 + r range reduction, a ~1 ULP degree-6 minimax on r, 2^k via integer exponent construction), with a scalar libm fallback for other targets, for CPUs without AVX2+FMA, and for the sub-8 tail of every row. Sharing one core is what keeps the two paths bit-identical: the max / exp / normalize arithmetic is literally the same instructions whether src == dst (in-place) or disjoint (out-of-place). Non-finite lanes are patched to match the scalar reference through the shared normalization: -inf and deep underflow flush to exactly 0, NaN stays NaN, so a fully masked row still normalizes to NaN and a partially masked row's masked positions are exactly 0. All existing softmax parity/bit-identity tests stay green (cargo test -p onnx-runtime-ep-cpu --lib: 1340 passed). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

The shipped CPU EP does not link MLAS, and its softmax evaluated one libm `f32::exp` per element. Once #1231 let the benchmark build the arm that actually ships, that showed up as 9.3-10.9x ORT across the whole softmax family - the single largest gap the CPU EP had. `scale_mask_softmax_serial` routes through the same helper, so fused attention paid it too. Both the in-place and out-of-place non-MLAS forms now go through one `softmax_row_core(src, dst, d)`, which evaluates exp eight lanes at a time with a Cephes-style range reduction: Cody-Waite split of ln2, a degree-6 minimax polynomial in FMA form, and 2^k built directly as (k+127)<<23. Dispatch is runtime `is_x86_feature_detected!`; the original scalar loop is retained for non-AVX2 targets and for sub-8 tails. Sharing one core is what keeps the two forms bit-identical. Non-finite lanes are handled by explicit masks rather than by hoping the polynomial behaves: NaN via an unordered compare plus blend, and -inf or deep underflow forced to exactly 0 via an ordered compare plus andnot. Both masks are falsifiable - disabling either one fails the suite. Measured with arms interleaved innermost, 8 reps, min-of-ratio (load average was ~10 from other work on the box, which inflates both arms): model before after speedup parity sm_prefill_h32_s512 9.044 2.005 4.51x PASS sm_bert_b8_s128 12.909 3.411 3.78x PASS sm_decode_h32_kv1024 11.067 1.412 7.84x PASS sm_decode_h32_kv2048 18.333 2.556 7.17x PASS sm_decode_h32_kv4096 31.286 4.286 7.30x PASS sm_decode_h32_kv8192 48.926 6.962 7.03x PASS Also marks the bit-identity test Miri-ignored, with the reason measured rather than assumed, and moves the #1231 feature guard into a const block so clippy is clean again. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Opus review found a blocking regression I introduced: moving the `mlas` feature guard into a `const` block made it a compile-time assert, and cargo features are additive and indistinguishable at `cfg` time, so it could not tell "mlas became a default" from "mlas was explicitly requested". It fired on both, and the CI lanes that build this crate with `--features mlas` - the MoE and QLinearMatMul reference-arm tests - then failed to compile the test target with E0080 rather than running. I had validated the mlas arm *before* making that edit and never re-ran it, which is exactly the ordering mistake the reviewer caught. Restored the runtime assert with an `allow`, and recorded in a comment why a const block is wrong here rather than merely lint-noisy. Also widens the bit-identity test from a single 8-wide row to 8, 11 and 6, so it covers a whole vector body, a body plus a 3-lane scalar tail, and a pure tail with no vector body. NaN in a vectorized lane is canonicalized by the mask while NaN in a tail lane is whatever libm returns, so a change that stopped both forms sharing one core would diverge on exactly these rows. The falsifier still fires at all three widths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…rt transpose (#1241) ## The defect The shipping (non-MLAS) MoE path allocated and filled a **full transposed copy of the expert weight matrix on every call, for every routed expert group**, purely because expert weights are stored `[out_features][in_features]` while `gemm` wants `[K][N]`: ```rust let mut weights_kn = vec![0.0f32; expert_stride]; for output_feature in 0..out_features { for input_feature in 0..in_features { weights_kn[input_feature * out_features + output_feature] = expert[output_feature * in_features + input_feature]; } } gemm(input_rows, &weights_kn, output_rows, group.len, in_features, out_features)?; ``` For the Mixtral shape that is a **29 MiB** allocate-and-transpose for fc1 (these graphs use `swiglu_fusion=1`, so `fc1_size = 2 * inter`, giving `7168 x 1024 x 4`) plus 14.7 MB for fc2 - **to multiply a single token row**. This was invisible until #1231, because `bench_generic` required the `mlas` feature and therefore only ever measured the reference arm. §18.3's claim that every MoE `t=1` cell sat at 1.006-1.052 was an MLAS-arm measurement. The arm that ships was **39-55x slower than ORT** at one token. ## The fix `gemm_bt(a, bt, c, m, k, n)` in `matmul.rs` computes `C[m,n] = A[m,k] * Bt[n,k]^T` directly on the `[out][in]` expert slice, so the transposed copy is never materialised. This layout is also strictly *better* for the kernel: for `C[i][j] = sum_k A[i][k] * Bt[j][k]` both operands are contiguous along `k`, so the inner loop is a pure contiguous dot product rather than a strided gather. AVX2+FMA with runtime `is_x86_feature_detected!` dispatch, following the structure of the `softmax_avx2` module merged in #1234: `loadu`/`storeu` so no alignment is assumed, `#[target_feature(enable = "avx2,fma")]` reachable only behind the runtime check, multiple independent accumulators to hide FMA latency, portable scalar fallback for non-x86 and for `k` tails. Parallelised with the *existing* `par_chunks_mut` / `for_each_init` / `worker_span` patterns already in `matmul.rs` - no new pool, no new scheduling policy (that is Sebastian's domain). Both non-MLAS sites in `moe.rs` (`grouped_linear` and `linear_grouped`) are rewired; the `weights_kn` allocation and its transpose loop are deleted outright. The MLAS reference arm is untouched. ### Cache blocking The first cut regressed the largest prefill shape (phi35moe t512: 2.18x, against 1.93x before). The tile reduces along `k` in the SIMD lanes, so a `bt` larger than L3 - the Phi experts are 52 MiB against 32 MiB of shared L3 - gets re-streamed from DRAM once per row band. Sweeping interior tiles **column-panel-first**, pinning a ~256 KiB `bt` slice in cache and reusing it across row bands, turned that regression into a 1.25x win and improved every other prefill cell too. ## Measured Production-native build, `nm -C | grep -ci mlas` = **0** on both arms. Thread-matched (`--native-threads 1 --ort-intra-threads 1`), arms interleaved innermost with the order alternated between reps, `native_min`/`ort_min`, 4 reps. Ratio is native/ORT, so lower is better and below 1.0 beats ORT. | model | before | after | speedup | parity | |---|--:|--:|--:|:--| | moe_mixtral_h1024_i3584_e8_t1 | 69.610 | **0.926** | 75.1x | PASS | | moe_qwen3moe_h2048_i768_e16_t1 | 41.381 | **0.912** | 45.4x | PASS | | moe_phi35moe_h2048_i6400_e4_t1 | 52.741 | 1.096 | 48.1x | PASS | | moe_mixtral_h1024_i3584_e8_t32 | 17.75 | **0.817** | 21.6x | PASS | | moe_qwen3moe_h2048_i768_e16_t32 | 8.08 | 1.027 | 8.0x | PASS | | moe_phi35moe_h2048_i6400_e4_t32 | 9.17 | 1.222 | 7.6x | PASS | | moe_mixtral_h1024_i3584_e8_t512 | 3.672 | 1.235 | 3.0x | PASS | | moe_qwen3moe_h2048_i768_e16_t512 | 1.88 | 1.390 | 1.3x | PASS | | moe_phi35moe_h2048_i6400_e4_t512 | 1.934 | 1.553 | 1.2x | PASS | The t=1, t=512 and mixtral rows are my own independent re-measurement; the t=32 rows are from the implementation run. Speedup is visible in both the min and the median of every cell, and no cell regresses. ## What still loses - stated plainly - **Decode t=1:** mixtral (0.926) and qwen3moe (0.912) now beat ORT. **phi35moe still loses ~10%** (1.096). That residual is bandwidth-bound rather than transpose-bound: phi35moe (i=6400) streams the most weight bytes per token of the three. - **Prefill t=512: every cell still loses** (1.24-1.55), though all improved and all now beat the old native path by a wide margin. This is microkernel efficiency against MLAS: the tile is NR=2 / 256-bit with no A/B packing, where MLAS uses wider register tiles and packed panels. Closing it means a packed, wider-NR prefill kernel. **MoE is not won.** It is no longer catastrophic. ## Tests Six new tests, and **every one was falsified** - the mutation used and the resulting failure: | test | mutation | observed failure | |---|---|---| | `gemm_bt_matches_naive_reference_across_shapes` | drop `acc3` in the final reduce | `index 0: got -0.10581067, want -0.06648377` | | `gemm_bt_handles_zero_extents` | remove the `m==0\|\|n==0` guard | panic `chunk_size must not be zero` | | `gemm_bt_block_scalar_matches_naive_reference` | stride `a` by `n` instead of `k` | `index 3: got 0.04070734, want -0.015676457` | | `gemm_bt_matches_transpose_then_gemm` | disable the `dot()` scalar k-tail | `index 0: gemm_bt 0, transpose+gemm -0.0043403152` | | `grouped_linear_matches_a_hand_rolled_transposed_gemm` | force expert 0 | `index 14: got 0.798, want 1.16` | | `gemm_bt_tile_kernel_matches_naive_single_threaded` | disable the `tile_4x2` scalar k-tail | `index 0: got -0.0057729427, want -0.009697894` | Coverage includes `m=0`, `n=0`, `k=0`, `m=1` (the decode case), `n=1`, `k` tails, and multi-panel `k` (512/520/300/259/577) so the column-blocked path and its tail are genuinely exercised. The MoE pin test uses transpose-then-`gemm` as its oracle, which is exactly the right reference here. - `cargo test -p onnx-runtime-ep-cpu --lib` - **1347 passed**, 0 failed. - `cargo test -p onnx-runtime-ep-cpu --lib --no-default-features --features mlas kernels::moe::` - **19 passed**; MLAS reference arm intact. - `cargo clippy -p onnx-runtime-ep-cpu --lib --all-targets` - clean. `rustfmt --check` clean on both changed files. - Parity PASS in all 9 models, both arms, every rep. ## Reported, not fixed `block_quantized_matmul.rs::dequantize_weight_kn` builds a `[k][n]` weight and feeds `gemm`, but it dequantizes *straight into* `[k][n]` rather than transposing, and caches the result for constant weights - amortised, not per-call, so it is not this antipattern. `conv_ref.rs` and `accelerate_gemm.rs` use the weight as the A operand. No other site in `kernels/` materialises a transpose into `gemm`. No ORT fallback, no defer, no runtime MLAS. ## Review Opus review returned **MERGE WITH FIXES** with **no blocking defects**. It independently reproduced four of the six falsifications verbatim, confirmed the `unsafe` is sound (every `#[target_feature]` entry is reachable only behind the runtime check, all loads are `loadu`, the parallel path hands out genuinely disjoint mutable slices), confirmed the MoE rewiring is faithful (`k=in_features`, `n=out_features`, offsets and activation/routing untouched), confirmed the MLAS arm still builds and passes, and reproduced the headline numbers (mixtral t1 0.935 against my 0.926; phi35 t512 1.549 against my 1.553). Two findings taken: 1. **Coverage gap.** The column blocking engages only at `n * k > 65536`, and every shape in the tile test stayed under it - the largest was `129x130x131` at 17030. The cache-blocking change, which is what turned the largest prefill shape from a regression into a win, was covered by nothing but the 629 MiB integration fixtures. Added `(8, 2000, 64)`, `(10, 2001, 70)` and `(1, 2000, 40)` (the multi-panel decode row). The new coverage is not decorative. Mutating the interior sweep so every row band after the first re-tiles the first panel's columns fails on `8x2000x64` with `got NaN, want -0.25505134`, while every pre-existing shape in the same test passes it. That is the gap, demonstrated. 2. **Wrong comment.** The large-`k` shapes were annotated as forcing "multi-panel accumulation in `tile_4x2`" with panels of `kc=44` and `kc=3`. There is no K blocking anywhere in `gemm_bt` - `tile_4x2` reduces all of `k` in a single pass, as its own doc says. Those shapes exercise a long single-pass reduction and its `k % 8` tail, and nothing panelled. Rewritten to say what they actually do. One finding **not** taken: the reviewer flagged "a 29 MiB strided scatter" in the `linear_grouped` doc as stale, computing 14.7 MB. That is the fc2 expert. These graphs are generated with `swiglu_fusion=1`, so `fc1_size = 2 * inter` and the Mixtral fc1 expert is `7168 x 1024 x 4` = 29 MiB. The comment was correct; the understatement was in this PR body, which I have corrected above. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…1242) Records Phase 11 and corrects the ratios that measured the wrong build. ## Why `bench_generic` declared `required-features = ["mlas"]` while containing no MLAS references at all. `mlas` is not a default feature of `onnx-runtime-ep-cpu`, so **every ratio in this document measured a configuration production never builds**. §29.4 spotted the harness problem and flagged it as open work; it underestimated the consequence. ## What it hid Two kernels had no vectorized production path at all, both masked for months by the MLAS arm covering for them: - **Softmax** called scalar `f32::exp()` once per element - **9.3-10.9x ORT**. `scale_mask_softmax_serial` shares the helper, so fused attention paid it too. Fixed in #1234. - **MoE** allocated and filled a full transposed copy of an expert's weight matrix *on every call, for every routed expert group* - 29 MiB for a Mixtral fc1 with `swiglu_fusion=1`, to multiply as little as **one token row**. **39-72x ORT at t=1**. Fixed in #1241. §18.3 had recorded every MoE `t=1` cell at 1.006-1.052 and concluded the only remaining MoE problem was thread scheduling at `t=8/16`. On the arm that ships, `t=1` was the worst cell in the entire benchmark suite. ## What this PR does Adds **§31** with: - the linker-level falsifier a feature flag cannot provide - `nm -C | grep -ci mlas` gives **0** on the shipping binary against **805** on the reference one; - the corrected production-native matrix for transforms and MoE (thread-matched, interleaved, min/median/max of per-rep ratios, parity PASS throughout); - **§31.4**, recording that the §19/§20 conclusion was wrong. Those sections closed the transpose question on the grounds that a layout-aware GEMM consumer was a cost "these graphs simply do not exercise" and would need a synthetic fixture built. The most expensive transpose in the runtime was never a `Transpose` node at all - it was inside the MoE kernel, feeding straight into a GEMM, on three of the models already in this document. - **§31.7**, two methodology rules: benchmark the artifact that ships and prove it with the linker rather than the flag; and treat an untested arm as unmeasured rather than "probably similar" - here it was 10x and 70x. - **§31.8**, where the CPU EP stands now, including what still loses: softmax 1.29-1.63x, MoE prefill 1.22-1.56x, phi35moe decode 1.10x, and the small-graph per-call overhead cells. Correction banners are added to §18.3, §27 and §29.4 rather than editing their tables, so the original numbers and the reason they misled both stay readable. Documentation only - no code changes. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

github-actions · 2026-08-18T13:22:53Z

🔴 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
🔴	`gather/large_bf16_threads=1-internal/131072`	10.25 µs	18.35 µs	+79.0%
🔴	`matmul/medium_generic_f32_threads=8/32x512x512`	928.00 µs	1.63 ms	+75.5%
🔴	`matmul/large_generic_f32_threads=8/32x1024x1024`	3.68 ms	6.21 ms	+68.5%
🔴	`block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024`	529.94 µs	852.05 µs	+60.8%
🔴	`matmul/large_generic_f16_threads=8/32x1024x1024`	79.17 µs	123.94 µs	+56.6%
🔴	`matmul/medium_generic_f16_threads=8/32x512x512`	27.69 µs	40.93 µs	+47.8%
🔴	`matmul/small_generic_f32_threads=8/1x256x256`	34.44 µs	50.39 µs	+46.3%
🔴	`grammar_masking/llguidance_compute_mask/32`	76.53 µs	109.86 µs	+43.6%
🔴	`gather/large_f16_threads=1-internal/131072`	11.91 µs	16.80 µs	+41.0%
🔴	`matmul/large_generic_bf16_threads=8/32x1024x1024`	1.25 ms	1.76 ms	+40.5%
🔴	`block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1`	432.32 µs	599.63 µs	+38.7%
🔴	`gather/large_f32_threads=1-internal/131072`	32.88 µs	44.69 µs	+35.9%
🔴	`matmul/small_generic_bf16_threads=8/1x256x256`	29.00 µs	39.27 µs	+35.4%
🔴	`matmul/medium_generic_bf16_threads=8/32x512x512`	351.53 µs	473.26 µs	+34.6%
⚠️	`matmul/small_generic_bf16_threads=1/1x256x256`	28.79 µs	36.59 µs	+27.1%
⚠️	`matmul/medium_generic_f16_threads=1/32x512x512`	27.46 µs	34.55 µs	+25.8%
⚠️	`matmul/large_generic_f32_threads=1/32x1024x1024`	8.93 ms	11.14 ms	+24.8%
⚠️	`add/medium_f32_threads=1-internal/262144`	26.83 µs	33.31 µs	+24.2%
⚠️	`matmul/large_generic_f16_threads=1/32x1024x1024`	93.41 µs	114.74 µs	+22.8%
⚠️	`logit_processing/seven_processor_chain_per_step`	335.65 µs	403.25 µs	+20.1%
⚠️	`matmul/small_generic_f16_threads=8/1x256x256`	31.39 µs	37.62 µs	+19.9%
⚠️	`matmul/medium_generic_bf16_threads=1/32x512x512`	496.05 µs	591.69 µs	+19.3%
⚠️	`matmul/small_generic_f16_threads=1/1x256x256`	28.96 µs	33.67 µs	+16.3%
⚠️	`block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024`	43.15 µs	50.15 µs	+16.2%
⚠️	`matmul/medium_generic_f32_threads=1/32x512x512`	2.16 ms	2.50 ms	+15.5%
✅	`matmul/small_generic_f32_threads=1/1x256x256`	36.74 µs	42.12 µs	+14.7%
✅	`matmul/large_generic_bf16_threads=1/32x1024x1024`	1.87 ms	2.14 ms	+14.6%
✅	`gather/small_f32_threads=1-internal/4096`	643.5 ns	722.5 ns	+12.3%
✅	`qwen3_sampling_processors/top_k_top_p_full_sort_baseline`	6.03 ms	6.65 ms	+10.3%
✅	`qwen3_sampling_processors/top_k_top_p_fast`	676.83 µs	734.42 µs	+8.5%
✅	`gather/small_bf16_threads=1-internal/4096`	479.2 ns	517.9 ns	+8.1%
✅	`block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024`	42.70 µs	45.92 µs	+7.6%
✅	`block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1`	155.08 µs	166.40 µs	+7.3%
✅	`reduce_mean/large_f32_threads=1-internal/262144`	973.18 µs	1.03 ms	+6.3%
✅	`kv_cache/alloc_dealloc_pages`	42.26 µs	44.31 µs	+4.9%
✅	`add/large_f16_threads=1-internal/4194304`	1.84 ms	1.91 ms	+4.1%
✅	`gather/small_f16_threads=1-internal/4096`	510.3 ns	528.8 ns	+3.6%
✅	`gather/medium_f32_threads=1-internal/32768`	4.12 µs	4.26 µs	+3.4%
✅	`qwen3_sampling_processors/top_p_fast_after_top_k`	524.58 µs	539.45 µs	+2.8%
✅	`tokenization/decode_tokens_per_second`	6.19 ms	6.36 ms	+2.7%
✅	`qwen3_sampling_processors/top_k_partial_selection`	161.06 µs	163.21 µs	+1.3%
✅	`gather/medium_bf16_threads=1-internal/32768`	2.41 µs	2.44 µs	+1.0%
✅	`reduce_mean/medium_f32_threads=1-internal/65536`	261.06 µs	259.46 µs	-0.6%
✅	`add/large_bf16_threads=1-internal/4194304`	1.94 ms	1.92 ms	-1.1%
✅	`sampling_latency/min_p_per_token`	218.77 µs	208.35 µs	-4.8%
✅	`qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline`	3.91 ms	3.67 ms	-6.3%
✅	`qwen3_sampling_processors/top_k_full_sort_baseline`	2.19 ms	2.01 ms	-8.5%
✅	`gather/medium_f16_threads=1-internal/32768`	2.89 µs	2.64 µs	-8.6%
✅	`tokenization/encode_tokens_per_second`	421.01 µs	381.22 µs	-9.5%
✅	`reduce_mean/small_f32_threads=1-internal/4096`	16.61 µs	14.84 µs	-10.7%
✅	`add/medium_f16_threads=1-internal/262144`	141.16 µs	124.78 µs	-11.6%
✅	`sampling_latency/greedy_per_token`	3.49 µs	3.04 µs	-12.9%
🟢	`add/small_f32_threads=1-internal/1024`	221.3 ns	186.7 ns	-15.6%
🟢	`add/medium_bf16_threads=1-internal/262144`	136.05 µs	110.90 µs	-18.5%
🟢	`add/small_f16_threads=1-internal/1024`	600.7 ns	489.0 ns	-18.6%
🟢	`sampling_latency/top_k_per_token`	64.10 µs	48.97 µs	-23.6%
🟢	`add/small_bf16_threads=1-internal/1024`	614.1 ns	466.3 ns	-24.1%
🟢	`add/large_f32_threads=1-internal/4194304`	1.07 ms	796.15 µs	-25.4%
🟢	`sampling_latency/top_p_per_token`	538.48 µs	366.02 µs	-32.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: { 5.34 4.01 6.24 }

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)

…1 ULP held) (#1245) ## What Replaces the Cephes `1 + r + r^2*P5(r)` form of `exp8` in `mod softmax_avx2` with a degree-6 Remez minimax polynomial evaluated in **pure Horner** form, whose two leading coefficients are pinned to exactly `1.0` so the `1 + r` head stays exact. Net effect per 8-lane `exp`: one fewer FMA-unit op and one fewer FADD (the separate `r*r` multiply and the trailing `+1` both disappear into the Horner chain). The Cody-Waite two-part `ln2` range reduction is **unchanged** - see "what I rejected". ## Why this, and why now #1234 vectorised this `exp` and took the softmax family from 9.3-10.9x ORT to 1.29-1.63x. Softmax is now the most-lost kernel family in the CPU EP. I had attributed that residual to the three-pass memory structure (max / exp+sum / normalize). **That hypothesis is falsified.** A two-pass online/streaming softmax measured *2.4x worse* across all 7 models, with the loss growing at the largest rows - the opposite of a bandwidth-bound signature. The streaming recurrence needs an `exp` for both the increment and the max-rescale, ~3 exp/element against the three-pass form's 1. That negative result is being written up separately. So the residual is **exp throughput**, and this PR attacks exactly that. It is a modest win, and I am not going to dress it up as more. ## Measurements - production-native only Build: `--no-default-features --features bench-native,cuda-13000`. Falsifier: `nm -C bench_generic | grep -ci mlas` = **0 for both arms** (binaries verified distinct by md5). Thread-matched `--native-threads 1 --ort-intra-threads 1`, arms interleaved innermost in one driver process with arm order alternating between reps. Ratio = native_min/ort_min; **lower is better, below 1.000 beats ORT**. Two independent 8-rep passes of my own (16 reps total), on top of the implementing agent's 12- and 14-rep passes. Medians: | model | before | after (pass 1) | after (pass 2) | |---|--:|--:|--:| | sm_bert_b8_s128 | 1.493 / 1.529 | **1.349** | **1.351** | | sm_decode_h32_kv1024 | 1.471 / 1.471 | 1.477 | **1.412** | | sm_decode_h32_kv2048 | 1.567 / 1.525 | **1.477** | **1.500** | | sm_decode_h32_kv4096 | 1.516 / 1.475 | **1.446** | **1.422** | | sm_decode_h32_kv8192 | 1.555 / 1.566 | **1.474** | **1.459** | | sm_prefill_h32_s512 | 1.398 / 1.400 | **1.300** | **1.311** | | sm_whisper_cross | 1.243 / 1.246 | **1.170** | **1.174** | `parity=PASS` on every cell of every rep. Pass 2 improves 7/7 on both min and median; pass 1 improves 7/7 on min and 6/7 on median (kv1024 median flat at 1.471 -> 1.477, which pass 2 resolves to 1.412). Roughly a **4-10% ratio improvement**, no model regressed. One honesty note on methodology: the *absolute* native times are not comparable across arms on this shared, bimodal box - one pass showed bert at 0.890 "native speedup" while its ratio vs ORT clearly improved, because contention moved under both arms. Since ORT is byte-identical code in both binaries and is timed in the *same process invocation* as native, the paired ratio is the internally-controlled metric, and it is the one I am quoting. ## Accuracy Exhaustive sweep over all 1.118e9 f32 values in the softmax domain `[-87.336, 0]` against a correctly-rounded f64 reference: - current `exp8`: **1 ULP** - this `exp8`: **1 ULP** No accuracy is spent to buy the speed. ## New test, and what it catches `exp8_stays_within_two_ulp_of_the_f64_reference` probes `exp8` *directly*, because the row-level tests structurally cannot police it: softmax divides every exponential by the row sum, so relative error shared across lanes **cancels in the ratio**. An `exp8` variant measured at 64 ULP against f64 still matched the scalar softmax to 2.4e-7. I falsified it with my own mutation, distinct from the implementing agent's: perturbing `c2` by ~6e-6 relative (`4.999_999_1e-1` -> `4.999_970_0e-1`). ``` exp8 exceeded the 2 ULP accuracy bar: max_ulp=6 at x=-0.32134318 ``` That mutation **fails only the new test** - all 21 pre-existing row-level softmax tests pass, including `out_of_place_softmax_is_bit_identical_on_pathological_rows`, `large_and_masked_logits_stay_finite` and `a_fully_masked_row_reproduces_the_reference_runtimes_nan`. So the test closes a real hole rather than restating existing coverage. ## What I rejected, and the measurement that decided it | candidate | decided by | outcome | |---|---|---| | magic-number rounding (replace `VROUNDPS`) | microbench ~0.98x | rejected - add+sub lengthens the `t -> k` chain | | `exp2` formulation (fast, 1-part log2e) | **64 ULP** | rejected - f32 `f = x*log2e - k` drops low bits | | `exp2` formulation (accurate, 2-part log2e) | 1 ULP but ~1.04x | rejected - reintroduces the 2 FMAs, smaller win | | ILP unroll x2 / x4 | monotonically slower | rejected - loop is FMA-throughput-bound, not latency-bound | | degree-5 pure Horner | **3 ULP** | rejected - exceeds the 2 ULP bar | The `exp2` result is the informative one: cutting the two Cody-Waite range-reduction FMAs is worth ~15%, but only by destroying accuracy; doing it *accurately* gives the FMAs back and erases the gain. **The Cody-Waite reduction is load-bearing** and I left it alone. The only accuracy-free lever left was the polynomial's shape. ## Semantics preserved - In-place and out-of-place remain bit-identical (they still share one core). - Non-finite handling untouched: fully-masked rows -> NaN, NaN -> NaN, `-inf` and deep underflow -> exactly `0`. - No ORT fallback, no defer, no runtime MLAS. `mlas` remains non-default. ## Checks - `cargo test -p onnx-runtime-ep-cpu --lib`: **1348 passed, 0 failed**, 16 ignored - `cargo clippy -p onnx-runtime-ep-cpu --lib`: clean - `cargo fmt -p onnx-runtime-ep-cpu -- --check`: clean - `cargo check -p onnx-runtime-ep-cpu --lib --features mlas`: **compiles** - explicitly re-verified on the final tree, because in #1234 I shipped a `const {}` assert that broke this exact lane after I had validated an earlier state. ## Standing This does not close the softmax gap - it is still 1.17-1.50x ORT. It removes ~4-10% of it without spending accuracy, and it narrows the remaining cause: with the memory- structure hypothesis falsified and the range reduction proven load-bearing, what is left is ORT's exp/vectorisation quality or per-call overhead. Still not complete. --------- Signed-off-by: Justin Chuby <justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Roy <roy@squad.local> Co-authored-by: resch <resch@users.noreply.github.com> Co-authored-by: Resch <resch@squad.local> Copilot-Session: 0190e2eb-abe4-451f-b36d-44a035a99b7e Copilot-Session: c80f8522-983c-47f7-8241-2155a823aabe

justinchuby and others added 4 commits August 18, 2026 07:19

merge main

d215f16

justinchuby marked this pull request as ready for review August 18, 2026 08:10

justinchuby merged commit 4d231ea into main Aug 18, 2026
7 of 17 checks passed

justinchuby deleted the squad/leon-p18-native-exp branch August 18, 2026 08:18

justinchuby mentioned this pull request Aug 18, 2026

perf(cpu-ep): layout-aware gemm_bt for MoE, killing the per-call expert transpose #1241

Merged

justinchuby mentioned this pull request Aug 18, 2026

docs(benchmarks): record Phase 11 and correct every MLAS-arm ratio #1242

Merged

justinchuby mentioned this pull request Aug 18, 2026

perf(cpu-ep): cheaper softmax exp8 via degree-6 pure-Horner minimax (1 ULP held) #1245

Merged

Navigation Menu

Repository navigation

perf(cpu-ep): vectorise the production softmax exp on AVX2+FMA - #1234

perf(cpu-ep): vectorise the production softmax exp on AVX2+FMA#1234
justinchuby merged 4 commits into
mainfrom
squad/leon-p18-native-exp

justinchuby commented Aug 18, 2026 •

edited

Loading

Uh oh!

Uh oh!

github-actions Bot commented Aug 18, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

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

Uh oh!

Why

What

Non-finite lanes

Measured

Miri

Validation

Review

Uh oh!

Uh oh!

github-actions Bot commented Aug 18, 2026

🔴 Benchmark Regression Detected

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

justinchuby commented Aug 18, 2026 •

edited

Loading