Repository navigation
perf(cpu-ep): unblock softmax's two bottlenecks (1.25x on the row kernel) - #1416
Merged
Merged
Conversation
Owner
Author
|
Status: validated locally, held on required CI — not merging. Per instruction, this waits for the required checks (
So the required checks are currently unreachable, and this PR stays open until they run. It is not blocked on review or on any known defect. Local validation on this branch (not a substitute for CI, recorded for whoever merges):
|
justinchuby
enabled auto-merge (squash)
August 19, 2026 06:41
Owner
Author
|
Status: auto-merge armed, waiting on required CI. Not merging by hand. GitHub Actions has not concluded a run on this repo for some time (last 200 runs: 0 #1429 has to land before this one can go green. While replicating the CI lanes locally I Independent replication of both required checks on
|
justinchuby
force-pushed
the
squad/leon-p28-softmax-phases
branch
2 times, most recently
from
August 19, 2026 17:15
55f6ea7 to
0fe8efc
Compare
justinchuby
added a commit
that referenced
this pull request
Aug 19, 2026
…x by phase (#1423) Ledger write-up for #1402 and #1416. No code. Records, in order: - **§40.1** The experiment §39.4 left open — one immutable pack shared across row blocks — **built, measured, and lost**, at every panel width and worse the wider the pool (geomean 0.20–1.04 vs the unpacked driver). Full grid included so it is not rediscovered. Two causes: a fork-join per panel, and a pack that is a partly serial Amdahl term (at phi35's `k=6400` a 512 KB panel is a single micro-panel, so splitting the pack by micro-panel leaves it on one thread). - **§40.2** The win one level up: distribute whole panels, which gets the "pack each panel exactly once" property §39.4 wanted, without a barrier and with the panel resident in the owning core's L2. - **§40.3** The gate, including the `k >= 1024` variant that scored better on grid geomean and was **rejected** for regressing a production shape. - **§40.4** Softmax profiled by phase: pass 2 is 73% of the row and throughput-bound, pass 1 is 11% and latency-bound, so the same fix helps one and hurts the other. Includes the magic-number round fusion that was measured, costed at 3%, and declined. - **§40.5** Two method lessons: exhaustive verification over all 2^32 f32 inputs (0.6s, and the control that proves it can fail), and the invariance-hidden test hole that let a deleted accumulator chain pass every existing test. Also records the mixtral cell where the §39 control arm moved +66% and the ratio metric had to be abandoned for native-time-vs-control. --- ## Update: adds §42 as well `main` claimed `## 40` while this branch was open, so my original section was renumbered to **§41** on rebase. This PR now also carries **§42**, recording the softmax fan-out gate finding from #1484. §42 is deliberately a continuation of upstream's §40 rather than a new theme. §40 closes by asking that "the fourth instance is recognised rather than re-derived"; §42 is that fourth instance, and it reports the way it *differs* from the three §40 collected: the constant was not calibrated in the wrong regime, it was expressed in the wrong unit. A row is not a unit of work — `n` prices work only if `d` is held fixed, and `d` is the key length. §40.3's remedy — record the regime beside the constant — would not have caught it. §42.3 records the more transferable half: softmax rows are independent, so the output is bit-identical whether or not the fan-out happens, which makes every correctness test in the file blind to the gate by construction. Inverting the caller's use of the predicate undoes the entire optimisation and left all 1448 tests green. A performance gate has no numerical signature and has to be asserted directly. §42.4 is a reporting note: the host was heavily contended for that run (A/A nulls up to 61%), and the honest read came from structure rather than the grid — only four of seven fixtures change gate decision at all, so the other three cannot have moved and their scatter calibrates the host. #1484's doc comment cites §42, so the two are written to land consistently in either order. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…nel)
Profiling the row kernel by phase, rather than as a whole, showed the work
is not where the shape of the code suggests. Per element at d=1024, on this
host:
pass 1 (max) 67 ps
pass 2 (exp+sum) 446 ps
pass 3 (normalize) 58 ps
whole row 608 ps
Pass 2 is 73% of the row and pass 1 is 11%, and each is limited by a
different thing. Both limits are removable, and neither removal changes a
single output bit.
Pass 1 is latency-bound. It is one `vmaxps` chain, and `vmaxps` has ~4
cycles of latency, so the chain retires 8 floats every 4 cycles while the
load ports could feed 16. Four independent chains hide the latency
completely: 67 -> 37 ps/element, 1.8x. `max` is associative and commutative
over the values that survive, and feeding the loaded vector as the *first*
`max_ps` operand keeps the existing NaN convention (a NaN lane returns the
accumulator) in each chain, so no accumulator can ever hold a NaN and the
final fold obeys the same rule. The result is bit-identical to one chain.
Pass 2 is throughput-bound on ports 0/1, not latency-bound: the loop issues
16 port-0/1 ops per 8 lanes, which at 2 ops/cycle is the 8 cycles/iteration
measured. Four independent sum chains, the obvious counterpart to the pass-1
fix, are therefore *slower* (446 -> 458 ps/element) and are not used.
Three of those 16 ops exist only to patch `exp8`'s non-finite lanes: a
`max_ps` clamp, an unordered compare, and a `blendv`. They are unnecessary,
because both non-finite cases fall out of the arithmetic:
* `-inf` poisons `r` into `NaN`, but `is_zero` is true for it, and the
underflow mask that already exists is a bitwise `andnot` -- it produces
the required exact `0` whatever the poisoned build contained.
* `NaN` survives `round`, both `fnmadd`s and the Horner chain; `cvtps_epi32`
maps it to the integer indefinite `0x8000_0000`, whose low nine bits
build `pow2` = `1.0`, so `y` = `NaN * 1.0` = `NaN`, and `is_zero` is
false for it because the compare is ordered.
Dropping all three: 446 -> 370 ps/element, 1.20x.
Together, 608 -> 481 ps/element at d=1024; 1.16x at d=64 rising to 1.26x at
d>=256. Every softmax consumer shares this one kernel, including
`FusedAttention` through `scale_mask_softmax_rows`.
Not taken: folding the `mul`+`round` into one FMA with a magic-number
constant saves a fourth port-0/1 op, but it is worth only 3% (370 -> 356
ps/element) and it moves `k` for arguments that sit on a rounding boundary,
which would forfeit bit-identity and require re-establishing the 1-ULP bound
from scratch. Not worth 3%.
Numerics
--------
The `exp8` change is verified exhaustively rather than argued. An offline
harness evaluates the old and new forms over **all 2^32 f32 bit patterns**:
non-NaN inputs whose result differs in any bit: 0
NaN inputs where either side is not NaN: 0
So the 1-ULP bound over `(-87.336544, 0]`, the exact `0` at and below the
underflow threshold, and the `NaN` -> `NaN` contract all transfer unchanged
by construction, not by re-measurement. (Control: deleting the underflow
mask makes the same harness report 1,020,022,810 mismatches, first at
`-87.33654`.) The only observable difference is the `NaN` *payload*, now the
quieted input rather than a canonical `f32::NAN` -- which the row contract,
"a poisoned row normalizes to NaN", does not distinguish.
Two tests are added because both properties are now emergent rather than
explicit:
* `exp8_maps_the_non_finite_lanes_without_an_explicit_patch` pins the
vector-level contract with quiet, signalling and both-signed NaNs and
with five flush-to-zero inputs, placed in every lane in turn, plus the
one-ULP-above-threshold boundary.
* `the_row_maximum_covers_every_position` closes a hole the pass-1 change
exposed: softmax is invariant to the value subtracted, so a row maximum
that misses part of the row is numerically silent and *every existing
test passed with one of the four chains deleted*. It only surfaces as
overflow, so the test puts a 3e38 logit at each position of each of
seven `d` values and requires the row to stay finite.
Falsified: dropping a max chain fails the new row test; dropping the
underflow mask fails the new exp8 test and `large_and_masked_logits_stay_finite`;
reordering the fold (a no-op) stays green. Skipping the 8-wide drain in
pass 1 also stays green, correctly -- the scalar tail still covers those
elements.
Measurements
------------
End to end over all seven softmax fixtures x 1/2/4/8 threads, both arms
proved pure-native (`nm -C | grep -ci mlas` = 0), `--null-control` on every
run. 25 of 28 cells show the old kernel slower by more than that cell's
noise floor, from +2.0% to +51.8%; no cell shows a regression. The three
that landed inside the floor were re-run at 41 trials: all three then favour
the new kernel (kv1024 t=8 +10.1% > 2.5% floor; whisper t=4 +2.2%, prefill
t=4 +3.8%, both inside their floors but with native-time medians 3.7% and
7.2% below the old arm and level with the null).
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Review caught: the pathological-rows test still described the canonical-NaN blend this change removes; the pass-1 note said `vmaxpd` where the code is `vmaxps`; and one measurement was quoted as 442 rather than 446 ps/element. Also record the one caveat to the associativity argument that a differential sweep over ~588k rows turned up: a row whose maximum is zero can pick +0.0 or -0.0 depending on the grouping. It cannot reach the output, because both signs subtract identically and exp8(±0.0) is 1.0, and the sweep confirms every row is bit-identical -- but the comment claimed more than it had. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
`parallel_rows_per_task` refused to fan out below 64 rows however much work those rows held. Decode attention is `heads x kv_len`, so a 32-head model ran single-threaded at every pool width -- 32x8192 is 256 Ki elements, 16x the element floor sitting next to the row floor in the same predicate. Measured against ORT on a pure-native default build, that is the whole of the 2.2x-5.6x multi-thread gap on the decode softmax fixtures; native time was flat to within 1.4% from one thread to eight. Drop the row floor and keep the element floor, which measures work in the unit work is done in. This admits nothing new per chunk: chunks are sized by ROW_TILE_BYTES, so the newly-admitted 32x8192 yields 32 chunks of 8192 elements, exactly what the already-admitted 64x256 yields -- the test asserts that equality rather than the arithmetic behind it. `n >= 2` replaces the floor and is an impossibility rather than a threshold: a chunk is a whole number of rows. Native time falls 29.9%-43.6% wherever the fan-out engages (kv1024/2048 from t=4, kv4096/8192 at t=8), and the ours/ORT ratio follows: -53.9% at kv4096 t=8, -44.1% at kv2048 t=8, -39.9% at kv8192 t=8. The three fixtures that already cleared the old floor take the same code path before and after, and are reported unchanged. Split the work predicate out as `fan_out_is_worthwhile` so the policy is testable without a multi-threaded runtime. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Review found two mutations of the gate that the whole suite tolerated. Inverting `parallel_rows_per_task`'s use of `fan_out_is_worthwhile` pins every large softmax back to a single thread -- undoing this change entirely -- and 1448 tests still passed. Nothing could catch it: rows are independent, so the output is bit-identical whether or not the fan-out happens, and the new unit test exercised the predicate in isolation rather than through its caller. `the_fan_out_gate_is_wired_the_right_way_round` closes that by asserting the Some/None decision itself. Its refusals need no pool, so they run everywhere; the fan-out half returns early below two lanes, matching the existing `parallel_output_rows_dispatches_to_the_task_runtime`. The element floor was also pinned only at MIN and MIN-2, so relaxing it by one element survived. 3x5461 is one element short and now straddles it exactly. Also drop a comment reference to MIN_PARALLEL_SOFTMAX_ROWS, deleted last commit. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Upstream landed a new section 41 while this branch was open, so the ledger section recording this measurement is 43 rather than 42. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
squad/leon-p28-softmax-phases
branch
from
August 19, 2026 20:39
0fe8efc to
9e1c5f5
Compare
🔴 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❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1416 +/- ##
===========================================
- Coverage 82.10% 80.21% -1.89%
===========================================
Files 12 377 +365
Lines 5471 165927 +160456
Branches 5471 165927 +160456
===========================================
+ Hits 4492 133103 +128611
- Misses 780 27987 +27207
- Partials 199 4837 +4638
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
justinchuby
added a commit
that referenced
this pull request
Aug 23, 2026
…d a knob that no longer exists (#1822) Follow-up to #1173, correcting two defects I shipped in it and repairing the rule they undermined. Docs, one ledger string, one new test, one new script. No production kernel or routing change. ## 1. The ledger named a route gate that had already been deleted `PLAN[MatMulF32].shape_gate` said the native `SimdX86` route "gates M=1 on `ONNX_GENAI_CPU_MM_SIMD_M1_GEMV` (default off, #1116)". #1183 shipped that GEMV on by default and removed the probe. `git merge-base --is-ancestor 5417d04 bdb4599` confirms it landed **before** #1173 merged — so the ledger was wrong the day it landed. Today `sgemm_simd` calls `sgemm_simd_variant(a, b, c, m, k, n, true)` unconditionally and `use_m1_gemv` is a plain parameter that only the in-process A/B harness passes as `false`. No environment variable reaches that route. `docs/performance/CPU_MATMUL_ASSIGNMENT.md:559` already recorded the correct fact ("It is measured now, and the route is the default. There is no env probe on the dispatch any more"). Two files in the same directory disagreed and nothing compared them. **Now guarded.** `ledger_prose_only_names_environment_variables_that_still_exist` requires every `NXRT_*` / `ONNX_GENAI_*` token in the ledger's prose to still exist as a string literal in the crate's sources. It cannot check that the description is *right*, only that the knob is *real* — which is the half that goes stale silently. Mutation-verified, not just observed green: ``` matmul_f32: ledger prose names environment variable `ONNX_GENAI_CPU_MM_SIMD_M1_GEMV`, but no source file in this crate contains the literal "ONNX_GENAI_CPU_MM_SIMD_M1_GEMV". ``` ## 2. The doc published a toggle A/B that could not have been run #1173 carried a table captioned **"same binary, same session, toggle the only difference"**, reporting `decode 1×2048×2048` at 0.146 with `ONNX_GENAI_CPU_MM_SIMD_M1_GEMV` off against 0.337 with it on, and called turning it on "the obvious next slice". Nothing reads that variable. Setting it measures the same route twice; it cannot produce two different columns. The table is withdrawn and the retraction kept in the text rather than quietly deleted. This is the failure mode the document's own graduation rule warns about — **an arm that was not on the route it was labelled with** — committed by the document that wrote the rule. It survived review because a plausible number in a well-formed table is not self-evidently unmeasured. Readers are pointed at `bench_f32_gemm_ab`, which holds the route as a function parameter and carries the M≥2 rows as a built-in control. ## 3. The gap table is re-measured and the ≥5% rule is repaired The old table was one unguarded invocation per row at an unstated width, taken before the decode-placement corrections (#1729, #1794, #1811) — i.e. when the decode pool put 16 workers on 8 physical cores. New harness: `scripts/bench_native_vs_mlas_width.py`. Arms interleaved rep by rep so host drift lands on both equally; per-rep `os.wait4` CPU-efficiency guard adapted from #1809; six reps per arm; two widths. Raw verdicts, spreads and discards are all reported rather than summarised away. **Three findings, all about method rather than kernels.** | | narrow (6 cores, 1 L3) | wide (32 logical CPUs) | |---|---|---| | `matmul_f32 16×512×512` | 1.581, spread 41% | 0.866, spread 134% | | `matmul_f32 decode 1×2048×2048` | 1.117, spread 21% | 0.934, spread 13% | - **Two cases change verdict on width alone.** Same binary, same half-hour, only the CPU mask differs. `x86_sgemm` parallelises over column strips and MLAS declines to parallelise some shapes, so interleaving the two *routes* inside one process does not protect the ratio — it changes both at once. - **`16×512×512` disagrees with itself on both arms**, alternating `keep-mlas` / `native-graduates` from a byte-identical binary. **One more run of the old table could have graduated a route on this row.** - **The narrow arm is more trustworthy despite having fewer cores** — spreads 4–42% against 5–134%, and it lost no reps to the guard. Isolation beat parallelism. **Softmax now decomposes cleanly**, because no vendored MLAS kernel has changed since #1173 (the only `mlas-sys` edits are the additive straggler handshake in `work_stealing_pool.rs`, #828/#1714, which adds waiting). At matched width the MLAS control arm is stationary to within 4% while native improved **1.24–1.27×** — matching #1416's claim for the row kernel. The f32 GEMM rows get no such attribution and now say so explicitly: their control moved **2.0× the wrong way**, so only the current ratio at a stated width is defensible. **The rule gains what it lacked**: spread must be smaller than the claimed win; reps that did not get the CPU are discarded rather than averaged; a verdict is valid only at a stated width. Under it, `decode 1×2048×2048` — the first f32 GEMM case to show a real native win — **still does not graduate**: it costs more CPU (cpu_ratio 0.875), does not hold at 32 threads, and its 21% spread exceeds its 12% win. ## The width claim is verified, not asserted #1815 landed while this was in progress and observed the neighbouring `bench_generic` harness spawning its ORT arm *outside* the affinity confinement it applied to the native arm. That hazard applies to any `taskset` claim, including mine, so I checked it instead of trusting it — sampling `Cpus_allowed_list` from `/proc/<pid>/task/*/status` 40× across a live narrow-arm run: ``` '16,20,22,26,28,30': 478 observations native_vs_mlas- 273, mlas-sys-ws-0..4 39 each, nxrt-task-0..4 2 each '0-31': 1 (the taskset process itself, before exec) ``` Both routes confined identically; no thread escaped. The rule now requires this check. ## Validation - `dispatch_ledger` **17/17**, including the new falsifier, after merging latest `main`. - `default_artifacts_are_mlas_free` **9/9** — the no-MLAS-in-defaults invariant is untouched. - `cargo clippy -p onnx-runtime-ep-cpu --lib --all-targets` clean; `cargo fmt --check` clean. - Normal merge of `origin/main` (`aee2b9d11`), no rebase, no conflicts. ## Limitations - The narrow arm is six cores on one L3 of one x86-64 host. Nothing here transfers to aarch64 or to a two-socket box. - The `activations erf 1 Mi` row shows native 13.5% slower at matched width. The nearest scatter figure is the wide arm's 8% spread, but that is a spread of *ratios* against a move in a *native time*, so the two are not strictly commensurable. Its MLAS control also moved 11%. **Flagged for pinned re-measurement, not reported as a regression.** - The wide arm was taken with ~4–5 cores of unrelated load present. That is stated in the doc rather than hidden, and it is why its spreads are wider; the guard reports which reps were discarded instead of pretending the host was quiet. - No production behaviour changes here, so there is no performance claim to make about the shipped artifact. Refs #1173, #1183, #1809, #1815, #1416. ## Independent review, and what it changed An independent adversarial review of the full diff returned **no blockers** — it confirmed the ancestry argument behind the retraction, the stationary-control premise for the softmax attribution, and that the headline case is correctly *refused* by the rule (21% spread against a 12% win). It also found seven real defects, all now fixed in `f0323f9ed`. The one that mattered most was in the new test. It only proved the variable name appeared *somewhere* in the crate, so a variable whose read site had been deleted but whose name survived in an `EnvVarGuard::set(...)` line would still have passed — which is the precise shape of the defect this PR exists to correct. The test now requires the matching line to be an `env::var(` / `env::var_os(` read or an `_ENV: &str =` binding. Verified by mutation in **both** directions: | mutation | before | after | |---|---|---| | reinsert retired `ONNX_GENAI_CPU_MM_SIMD_M1_GEMV` into ledger prose | fails ✅ | fails ✅ | | retire the two real `NXRT_CPU_GEMM_BACKEND` reads, leaving the literal only in test guards | **passes ❌** | fails ✅ | The remaining six were prose defects in the doc: a stated spread range that contradicted its own table's 82% row, "within 4%" against a table reading −4.2%, a narrow-arm ratio fused with a wide-arm attribution, a spread quoted as 7.5% that was 8% *and* compared against an incommensurable quantity, the CPU-efficiency guard oversold as "what makes this table measurable at all" (in-process interleaving is what protects the ratio; the guard catches only *differential* descheduling), and a one-directional provenance argument standing in for the direct control measurement that actually carries the softmax attribution. **Two further defects I found myself while checking the tables against each other**, neither raised by the review: - The `ratio` column is a median of per-rep ratios while the `ns/unit` columns are medians of times. Medians do not distribute over division, so every row looked internally inconsistent to anyone who tried to divide it out (`0.0684 / 0.0617 = 1.109` against a stated `1.117`). Now documented, along with why the per-rep form is the correct one to quote: it pairs each MLAS invocation with the native invocation it was interleaved against, which is the entire point of interleaving. The then→now figures are relabelled as quotients of medians. - "wider than nine of the twelve wide-arm rows" was eleven of twelve. ## Adopting #1814 `aee2b9d11` (#1814) landed on `main` while this was in review, and it closes the exact hole the review found in the guard this document recommends. A differential CPU-efficiency check cannot see contention that lands evenly on both arms; #1814's confined-set meter reads busy jiffies on the process's own `Cpus_allowed_list` and subtracts the process's own CPU, so foreign load shows up directly. The rule now points at it, and the tables here are explicitly marked as predating it and guarded by the weaker method. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Profiling the row kernel by phase, rather than as a whole, showed the work
is not where the shape of the code suggests. Per element at d=1024, on this
host:
Pass 2 is 73% of the row and pass 1 is 11%, and each is limited by a
different thing. Both limits are removable, and neither removal changes a
single output bit.
Pass 1 is latency-bound. It is one
vmaxpschain, andvmaxpshas ~4cycles of latency, so the chain retires 8 floats every 4 cycles while the
load ports could feed 16. Four independent chains hide the latency
completely: 67 -> 37 ps/element, 1.8x.
maxis associative and commutativeover the values that survive, and feeding the loaded vector as the first
max_psoperand keeps the existing NaN convention (a NaN lane returns theaccumulator) in each chain, so no accumulator can ever hold a NaN and the
final fold obeys the same rule. The result is bit-identical to one chain.
Pass 2 is throughput-bound on ports 0/1, not latency-bound: the loop issues
16 port-0/1 ops per 8 lanes, which at 2 ops/cycle is the 8 cycles/iteration
measured. Four independent sum chains, the obvious counterpart to the pass-1
fix, are therefore slower (446 -> 458 ps/element) and are not used.
Three of those 16 ops exist only to patch
exp8's non-finite lanes: amax_psclamp, an unordered compare, and ablendv. They are unnecessary,because both non-finite cases fall out of the arithmetic:
-infpoisonsrintoNaN, butis_zerois true for it, and theunderflow mask that already exists is a bitwise
andnot-- it producesthe required exact
0whatever the poisoned build contained.NaNsurvivesround, bothfnmadds and the Horner chain;cvtps_epi32maps it to the integer indefinite
0x8000_0000, whose low nine bitsbuild
pow2=1.0, soy=NaN * 1.0=NaN, andis_zeroisfalse for it because the compare is ordered.
Dropping all three: 446 -> 370 ps/element, 1.20x.
Together, 608 -> 481 ps/element at d=1024; 1.16x at d=64 rising to 1.26x at
d>=256. Every softmax consumer shares this one kernel, including
FusedAttentionthroughscale_mask_softmax_rows.Not taken: folding the
mul+roundinto one FMA with a magic-numberconstant saves a fourth port-0/1 op, but it is worth only 3% (370 -> 356
ps/element) and it moves
kfor arguments that sit on a rounding boundary,which would forfeit bit-identity and require re-establishing the 1-ULP bound
from scratch. Not worth 3%.
Numerics
The
exp8change is verified exhaustively rather than argued. An offlineharness evaluates the old and new forms over all 2^32 f32 bit patterns:
So the 1-ULP bound over
(-87.336544, 0], the exact0at and below theunderflow threshold, and the
NaN->NaNcontract all transfer unchangedby construction, not by re-measurement. (Control: deleting the underflow
mask makes the same harness report 1,020,022,810 mismatches, first at
-87.33654.) The only observable difference is theNaNpayload, now thequieted input rather than a canonical
f32::NAN-- which the row contract,"a poisoned row normalizes to NaN", does not distinguish.
Two tests are added because both properties are now emergent rather than
explicit:
exp8_maps_the_non_finite_lanes_without_an_explicit_patchpins thevector-level contract with quiet, signalling and both-signed NaNs and
with five flush-to-zero inputs, placed in every lane in turn, plus the
one-ULP-above-threshold boundary.
the_row_maximum_covers_every_positioncloses a hole the pass-1 changeexposed: softmax is invariant to the value subtracted, so a row maximum
that misses part of the row is numerically silent and every existing
test passed with one of the four chains deleted. It only surfaces as
overflow, so the test puts a 3e38 logit at each position of each of
seven
dvalues and requires the row to stay finite.Falsified: dropping a max chain fails the new row test; dropping the
underflow mask fails the new exp8 test and
large_and_masked_logits_stay_finite;reordering the fold (a no-op) stays green. Skipping the 8-wide drain in
pass 1 also stays green, correctly -- the scalar tail still covers those
elements.
Measurements
Re-measured from a production default build at the request of the current
task:
cargo build --release -p onnx-genai-bench --features bench-native, i.e.default features on,
mlasnever named. No historical MLAS-linked number isreused; every figure below comes from this run.
The artifact is proved MLAS-free four independent ways, for both arms:
nm -C | grep -ci mlasnm -D | grep -ci mlasldd | grep -ci mlasstrings -a | grep -c 'mlas-sys|MlasGemm|mlas_sys'bench_genericmeasures ours and ORT inside one invocation, alternatingwhich runs first, so the
ours/ORTratio is drift-free by construction;ab.pyadditionally reverses arm order every other trial.--null-controlruns the baseline binary a second time under another name, so every cell
carries its own A/A floor measured in the same invocation. 7 fixtures x
1/2/4/8 threads x 15 trials x 3 arms = 1260 invocations, 40 runs / 15 warmups
each.
sm_bert_b8_s128sm_bert_b8_s128sm_bert_b8_s128sm_bert_b8_s128sm_decode_h32_kv1024sm_decode_h32_kv1024sm_decode_h32_kv1024sm_decode_h32_kv1024sm_decode_h32_kv2048sm_decode_h32_kv2048sm_decode_h32_kv2048sm_decode_h32_kv2048sm_decode_h32_kv4096sm_decode_h32_kv4096sm_decode_h32_kv4096sm_decode_h32_kv4096sm_decode_h32_kv8192sm_decode_h32_kv8192sm_decode_h32_kv8192sm_decode_h32_kv8192sm_prefill_h32_s512sm_prefill_h32_s512sm_prefill_h32_s512sm_prefill_h32_s512sm_whisper_crosssm_whisper_crosssm_whisper_crosssm_whisper_cross26 improved, 2 within noise, 0 regressed. The two within-noise
cells are reported as measured and are not claimed:
sm_prefill_h32_s512t=4is +3.29% against a 6.18% A/A floor, and t=8 is +2.06% against a 41.08%
floor -- that cell is not measurable on this host at eight threads, in either
direction.
The headline the task asked about: at t=1, where the A/A floor is 0.00-0.76%,
the ours/ORT band moves from 1.114-1.329 to 1.054-1.197.
What this does not fix, and where the rest of the gap is
The same run localises the remaining multi-thread gap, and it is not in the
kernel this PR touches. Native p50 for this PR, by thread count:
sm_bert_b8_s128sm_decode_h32_kv1024sm_decode_h32_kv2048sm_decode_h32_kv4096sm_decode_h32_kv8192sm_prefill_h32_s512sm_whisper_crossThe five decode fixtures are flat to within 1.4% from one thread to eight.
They are
n=32rows xd=1024..8192, andparallel_rows_per_taskrefuses tofan out below
MIN_PARALLEL_SOFTMAX_ROWS = 64rows, regardless of how muchwork those rows hold -- 32x8192 is 256 Ki elements, 16x the element floor that
sits right next to it. ORT does split them, which is the whole of the 2.2x-5.6x
ratio at t>=2; it is not a kernel-quality gap and no amount of per-element work
removes it. Chasing it belongs in a separate PR against the gate, not here.
Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com
Update on merge: this PR was consolidated, and the local gate matrix it passed
What changed since the body above was written
The body above closes by saying the fan-out gate "belongs in a separate PR". That
PR was #1484 -- and it touched the same file, so under the "consolidate overlap
first" directive its three commits were cherry-picked onto this branch and #1484
was closed. This PR is therefore all three softmax changes, validated as one
unit:
exp8-- clamp/compare/blend removed, singleandnotunderflow mask (pure Horner).MIN_PARALLEL_SOFTMAX_ROWSdeleted, replaced by a purefan_out_is_worthwhile(n, d)that prices work (n * d), not row count. Thisis what unblocks the five flat decode fixtures the body above had written off.
Ledger sections §42 (phases) and §43 (the fan-out gate) landed separately in
#1423 (
d08b1909e).softmax.rs's doc comment cites section 43; verified stillcorrect against merged
main(upstream owns §40/§41, and stops there).Final 28-cell softmax matrix --
mainvs this branch, ours/ORT ratioInterleaved arms, A/A null control, quiet host, production default build (no MLAS).
sm_decode_h32_kv1024sm_decode_h32_kv2048sm_decode_h32_kv4096sm_decode_h32_kv8192sm_bert_b8_s128sm_prefill_h32_s512sm_whisper_cross27 cells improved, 1 noise, 0 regressed.
kv8192t=4 −65.5%,kv8192t=8 −61.0%,kv2048t=4 −60.8%.1.249–1.375 → 1.033–1.180. That is the exp8 + pass-1 work.
the old body called permanently flat now scale.
whispert=1 (0.998),prefillt=4 (0.939),kv1024t=8 (1.005, parity).Declared not-a-win, for honesty:
whispert=81.755 → 1.740is noise(−0.9%, inside the null control's spread). It is counted as noise above, not as
an improvement. No cell regressed.
kv4096t=2/t=4/t=8 improve but remain theworst residual ratios in the grid (1.9–3.2x) -- that fixture is still ORT-favourable
and is not claimed as solved.
Numerics: exhaustive, not sampled
exp8changed shape, so the contract is proven over the entiref32domainrather than spot-checked:
1,020,022,810 mismatches, first at
-87.33654. The harness is live, notvacuously passing.
Local gate matrix (run on
main+ all three of my PRs, then re-run on mergedmain)cargo fmt --all --checkcargo clippy --locked --all-targets -D warningsscripts/check_cross_compile.sh)clippy --target aarch64-unknown-linux-gnu --all-targets -D warningstask_runtime/strided/provider/dtype)--no-default-features --features mlas)nm -C,nm -D,ldd,strings= 0 on both bench binaries;wheel_packagingguard 2/2Windows ARM64 could not be executed on this host.
ort-sys'sbuild.rsrunsbindgen against the downloaded ORT headers, and clang cannot find MSVC's
stdlib.h; the crate chainonnx-runtime-ep-cpu → onnx-runtime-ep-api → onnx-genai-ort-sysis mandatory, so the EP crate cannot be isolated from it.Mitigation, and the reason this is safe: this diff introduces zero
target_ospredicates. Every
cfgit adds istarget_arch = "x86"/"x86_64",feature = "mlas", orcfg(test). Windows ARM64 adds no configuration dimensionthat these changes span, and the aarch64
target_archdimension is coveredabove. CI additionally has a native
windows-11-armjob.Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com