[justinchuby]

What

The f16/bf16 activation path was a three-pass sandwich: widen the input into an
f32 scratch the size of the whole tensor, run the kernel over all of it,
then narrow all of it back. On a 1 Mi tensor that is a 4 MB buffer streamed out
to memory and read back for no reason — the intermediate is dead the moment the
narrow has seen it.

This fuses the last two passes. When the output is a contiguous f16/bf16 tensor
and the closure declares no extra reads, write_mapped_reading evaluates the
kernel ~8192 elements at a time into a 32 KiB scratch and narrows each chunk
straight into the output, so the intermediate never leaves L1.

The code already carried the diagnosis and named this as the fix:

The fix for f16/bf16 is to fuse widen, compute and narrow into a single pass
per chunk so each thread keeps its slice in L2 — a different change, and one
that has to be measured on its own.

This is that change, measured on its own.

Why the chunking is bit-identical

Two separate things have to hold, and the second one is subtle enough that the
first version of this PR got it wrong (see the review section).

1. The closure must be a pure elementwise map. A closure that reads across
lanes would see a different window per chunk. The gate is structural rather than
a convention:

• the row-broadcast callers (BiasGelu, FastGelu with a bias) must
  declare their bias range through also_read, because that is what lets
  output_direct_write_eligible reject an aliasing output;
• a non-empty also_read is exactly what keeps them on the whole-tensor path.

A caller cannot reach the chunked arm without having already asserted it reads
nothing but x. All five write_mapped (empty also_read) call sites were
checked and are plain elementwise maps: Gelu ×2, QuickGelu, FastGelu
without bias, and the Tanh/Sqrt/Erf arm of unary elementwise.

2. No chunk may be shorter than SIMD_MIN_LEN. dispatch! picks scalar or
vector from the length of the slice it is handed, and for the approximated
kernels the two are not bit-identical. A plain chunks(NARROW_CHUNK) hands
the last chunk whatever is left over, so a tensor of NARROW_CHUNK + 16
elements would run its final 16 scalar where one pass runs them vectorised.
chunk_end absorbs a trailing remainder below SIMD_MIN_LEN into the chunk
before it, and keeps every other boundary on an 8-wide lane so the vector run
lands where one pass would have put it.

This matters because the module treats bit reproducibility as a correctness
property rather than a tolerance — see the host_chunk_len docs and the
thread-invariance tests — and bf16 results cross to the CUDA EP.

unsafe

The &mut [u16] is built under validate(), is_contiguous(), numel() == n,
and a dtype match on Float16/BFloat16 (which is what makes the element width
16 bits). Disjointness from x does not rest on access ordering — that
would not make overlapping &[f32]/&mut [u16] sound anyway. It rests on
to_dense_f32_widen always allocating a fresh buffer when it widens, so a
16-bit output cannot be aliased by its own widened input. A debug_assert
is the tripwire for a future caller that breaks that.

Measurement

Production A/B, plugin_path_ab_vs_plain_ort, 1 Mi elements, 1 thread,
60 iterations, box gated to ≥93% idle. Ratio is ours/ORT, lower is better.

case	ratio_p50 before	after	ours p50 before → after
bench_tanh_f16_1m	1.408	1.213	0.752 → 0.653 ms (−13.2%)
bench_exp_f16_1m	1.206	1.206	0.718 → 0.718 ms (flat)

Repeatability of the headline case — seven independent rounds, the last three
taken after the correctness fix, which did not move it:

round	ratio_p50	ours p50	ORT p50 (control)
1	1.217	0.6458	0.5306
2	1.214	0.6542	0.5389
3	1.214	0.6525	0.5376
4	1.212	0.6536	0.5391
5 (post-fix)	1.212	0.6562	0.5415
6 (post-fix)	1.214	0.6524	0.5372
7 (post-fix)	1.214	0.6590	0.5427

The ORT arm moves 1.0% across the whole set against a 13.2% effect, so the
result is an order of magnitude clear of its own control.

Exp does not move because it never reaches write_mapped — only Tanh,
Sqrt and Erf route through the f32 path today. Widening that routing is the
obvious follow-up and is deliberately not in this PR.

Small tensors are unaffected: the scratch is (NARROW_CHUNK + SIMD_MIN_LEN).min(n),
so for any n below a chunk it allocates exactly the same n floats the old
path did.

Context: why f16 was the thing to fix

A full 1 Mi sweep of the activation family at 1 thread showed the native kernels
are already ahead of ORT on 17 of 18 f32 cases — mish 0.101, celu 0.135,
elu 0.137, selu 0.145, log 0.203, sqrt 0.309, erf 0.428, tanh 0.444,
sigmoid 0.451 — with relu (1.026) the only exception, and that one is
already addressed by #1355. The two 16-bit cases were the real remaining
kernel-side losses, which is what pointed here.

Tests

• a_short_tail_does_not_fall_off_the_vector_path — the one that guards the
  seam. It searches at run time for an input where the scalar and vector
  paths disagree after narrowing to f16 (most f32 disagreements are washed out
  by the narrow, so checking in the f32 domain proves nothing), asserts a
  witness was found, and then pins fused == unfused for tails of 1, 8, 16 and 31.
• chunk_end_never_emits_a_short_trailing_chunk — pins the decomposition
  itself across ~250 tensor sizes around chunk boundaries: every chunk is either
  ≥ SIMD_MIN_LEN or the whole tensor, every interior boundary is lane-aligned,
  and no chunk overruns the scratch.
• fused_narrow_is_bit_identical_across_chunk_boundaries — f16 and bf16
  against a single-shot reference across Tanh/Erf/Sigmoid for seven tail
  lengths.
• a_row_broadcast_closure_stays_off_the_chunked_arm — a BiasGelu with row
  width 61 (never divides a chunk) over a multi-chunk tensor.

Every one was verified non-vacuous by mutating production code and observing the
failure. Reverting chunk_end to the naive form fails the seam test at index
8192 with f16 bits 5146 vs 5145.

Suites, on this branch:

• onnx-runtime-ep-cpu --lib, debug profile: 1451 passed, 0 failed
• plugin_ort_e2e with NXRT_REQUIRE_ORT_TESTS=1: 55 passed, 0 failed,
  including every_assigned_node_is_also_executed_by_this_ep — assigned ==
  executed conformance holds.
• cargo clippy -p onnx-runtime-ep-cpu --lib --all-targets: clean.
• cargo fmt -p onnx-runtime-ep-cpu -- --check: clean.

Review

Independently reviewed by Claude Opus 4.8 with a brief to attack the safety
argument and verify every test by mutation. Verdict: REQUEST CHANGES, one
blocker, which is the SIMD_MIN_LEN tail divergence described above. The
reviewer reproduced it with a Sigmoid f16 tensor of 8208 elements (bits 3455
vs 3454) and correctly observed that the original 137-element tail could not
catch it. Fixed in 887c22e51, along with the reviewer's point that the
aliasing justification was argued from access ordering rather than from the
allocation that actually guarantees it.

The reviewer also confirmed: all eight call sites audited and the structural
also_read gate holds; the FnOnce→Fn relaxation is safe (every closure
captures only Copy state, no Cell/RNG/lazy-init); the guard set matches
write_dense_f32_narrow; n == 0, numel() != n and non-x86 all fall through
unchanged; and small-n allocation is not a regression.

Scope

No ORT CPU fallback, no runtime or default MLAS — pure native x86, one function
plus its tests.

---

Re-validated on latest main (e0aedd0fa), 2026-08-19

Latest main merged in normally (no rebase). Re-measured paired against main
from the same box, 1 thread pinned, 60 iters / 10 warmup, ours_p50 / ort_p50:

arm	reps	ours p50 (ms)	ratio
main	3	0.7367 / 0.7375 / 0.7319	1.361 / 1.367 / 1.371
this PR	5	0.6536 / 0.6612 / 0.6577 / 0.6532 / 0.6506	1.213 / 1.226 / 1.226 / 1.218 / 1.217

Ratio 1.367 → 1.218 (median), our own time −11%. ORT drifted under 1%
between the arms (0.534–0.541 on main, 0.535–0.539 here), which is the
control. The distributions do not overlap.

Cross-platform fix (f9a97e853)

The fused arm is behind #[cfg(x86)], but NarrowFn, chunk_end,
NARROW_CHUNK, the test helper scalar_vector_divergence and the four tests
that exercise chunk boundaries were not. An aarch64 build warned on every
one of them, and CI builds with -D warnings — so this branch would have
failed cross-platform CI while a local x86 clippy run stayed green.

Now gated alongside the code they belong to: 0 warnings on both x86-64 and
aarch64. All four tests assert properties of the chunked arm, which does not
exist off x86, so gating them loses no coverage.

Validation on the merged base

• cargo fmt clean; cargo clippy --all-targets -D warnings clean
• 1556 onnx-runtime-ep-cpu tests, debug profile (debug_asserts live)
• 55 plugin conformance tests (NXRT_REQUIRE_ORT_TESTS=1, release)
• every_assigned_node_is_also_executed_by_this_ep and
  every_fixture_loads_with_cpu_fallback_disabled green — nothing defers,
  nothing falls back to the ORT CPU EP
• aarch64-unknown-linux-gnu --all-targets check clean, 0 warnings

[github-actions]

🔴 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
🔴	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	61.19 µs	144.27 µs	+135.8%
🔴	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	567.08 µs	1.28 ms	+125.8%
🔴	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	53.03 µs	116.66 µs	+120.0%
🔴	matmul/large_generic_bf16_threads=8/32x1024x1024	1.39 ms	2.40 ms	+72.4%
🔴	matmul/large_generic_f32_threads=8/32x1024x1024	3.84 ms	6.49 ms	+68.9%
🔴	matmul/medium_generic_f32_threads=8/32x512x512	1.01 ms	1.64 ms	+62.6%
🔴	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	402.82 µs	624.97 µs	+55.2%
🔴	matmul/small_generic_f32_threads=8/1x256x256	37.94 µs	58.14 µs	+53.2%
🔴	matmul/medium_generic_f16_threads=8/32x512x512	32.67 µs	48.47 µs	+48.3%
🔴	matmul/medium_generic_bf16_threads=8/32x512x512	402.78 µs	576.22 µs	+43.1%
🔴	gather/large_f32_threads=1-internal/131072	29.76 µs	41.44 µs	+39.3%
🔴	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	70.42 µs	97.96 µs	+39.1%
🔴	matmul/large_generic_f16_threads=8/32x1024x1024	87.07 µs	117.72 µs	+35.2%
🔴	add/medium_f16_threads=1-internal/262144	105.81 µs	141.31 µs	+33.5%
🔴	add/medium_f32_threads=1-internal/262144	26.04 µs	34.64 µs	+33.1%
⚠️	matmul/large_generic_f32_threads=1/32x1024x1024	9.42 ms	12.06 ms	+28.1%
⚠️	matmul/medium_generic_f16_threads=1/32x512x512	32.96 µs	42.02 µs	+27.5%
⚠️	add/medium_bf16_threads=1-internal/262144	105.13 µs	131.92 µs	+25.5%
⚠️	matmul/small_generic_f32_threads=1/1x256x256	39.18 µs	49.11 µs	+25.3%
⚠️	matmul/small_generic_bf16_threads=8/1x256x256	34.06 µs	42.52 µs	+24.9%
⚠️	matmul/small_generic_bf16_threads=1/1x256x256	33.50 µs	41.35 µs	+23.4%
⚠️	matmul/large_generic_bf16_threads=1/32x1024x1024	2.02 ms	2.49 ms	+23.4%
⚠️	gather/large_bf16_threads=1-internal/131072	11.61 µs	14.29 µs	+23.1%
⚠️	tokenization/encode_tokens_per_second	384.14 µs	461.87 µs	+20.2%
⚠️	matmul/small_generic_f16_threads=1/1x256x256	32.83 µs	38.93 µs	+18.6%
⚠️	gather/medium_bf16_threads=1-internal/32768	2.45 µs	2.90 µs	+18.2%
⚠️	add/large_bf16_threads=1-internal/4194304	1.69 ms	1.97 ms	+16.7%
⚠️	matmul/large_generic_f16_threads=1/32x1024x1024	80.96 µs	94.05 µs	+16.2%
⚠️	add/large_f32_threads=1-internal/4194304	675.28 µs	782.38 µs	+15.9%
⚠️	gather/medium_f32_threads=1-internal/32768	4.06 µs	4.70 µs	+15.8%
✅	gather/large_f16_threads=1-internal/131072	11.49 µs	13.18 µs	+14.7%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.39 ms	2.74 ms	+14.7%
✅	reduce_mean/medium_f32_threads=1-internal/65536	250.20 µs	281.64 µs	+12.6%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.01 ms	1.12 ms	+11.2%
✅	gather/small_bf16_threads=1-internal/4096	477.3 ns	529.5 ns	+10.9%
✅	gather/medium_f16_threads=1-internal/32768	2.45 µs	2.71 µs	+10.3%
✅	add/large_f16_threads=1-internal/4194304	1.67 ms	1.85 ms	+10.3%
✅	logit_processing/seven_processor_chain_per_step	327.76 µs	360.93 µs	+10.1%
✅	add/small_bf16_threads=1-internal/1024	487.3 ns	536.4 ns	+10.1%
✅	reduce_mean/small_f32_threads=1-internal/4096	15.33 µs	16.80 µs	+9.6%
✅	grammar_masking/llguidance_compute_mask/32	83.25 µs	90.40 µs	+8.6%
✅	matmul/small_generic_f16_threads=8/1x256x256	38.27 µs	41.29 µs	+7.9%
✅	gather/small_f16_threads=1-internal/4096	498.3 ns	535.5 ns	+7.5%
✅	qwen3_sampling_processors/top_k_partial_selection	153.53 µs	164.63 µs	+7.2%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	547.02 µs	581.65 µs	+6.3%
✅	kv_cache/alloc_dealloc_pages	41.39 µs	43.76 µs	+5.7%
✅	add/small_f16_threads=1-internal/1024	550.4 ns	577.6 ns	+4.9%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.35 ms	6.52 ms	+2.6%
✅	tokenization/decode_tokens_per_second	7.37 ms	7.45 ms	+1.1%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.39 ms	2.34 ms	-2.0%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	584.50 µs	562.69 µs	-3.7%
✅	sampling_latency/min_p_per_token	240.62 µs	230.98 µs	-4.0%
✅	qwen3_sampling_processors/top_k_top_p_fast	727.95 µs	697.76 µs	-4.1%
✅	gather/small_f32_threads=1-internal/4096	781.9 ns	747.5 ns	-4.4%
✅	add/small_f32_threads=1-internal/1024	259.2 ns	243.1 ns	-6.2%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	4.35 ms	4.03 ms	-7.5%
🟢	sampling_latency/greedy_per_token	4.26 µs	3.55 µs	-16.6%
🟢	sampling_latency/top_k_per_token	74.24 µs	59.62 µs	-19.7%
🟢	sampling_latency/top_p_per_token	565.18 µs	428.55 µs	-24.2%

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

Host info

CPU: Apple M1 (Virtual)
Cores: 3
OS: Darwin 25.5.0 arm64
Rust: rustc 1.97.1 (8bab26f4f 2026-07-14)
Load avg: { 3.07 3.09 5.05 }

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)

[codecov]

Codecov Report

❌ Patch coverage is 96.68874% with 5 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.39%. Comparing base (daada2f) to head (f9a97e8).
⚠️ Report is 80 commits behind head on main.

Files with missing lines	Patch %	Lines
...nnx-runtime-ep-cpu/src/kernels/simd_activations.rs	96.68%	4 Missing and 1 partial ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #1376       +/-   ##
===========================================
- Coverage   82.60%   80.39%    -2.21%     
===========================================
  Files          12      378      +366     
  Lines        5472   167927   +162455     
  Branches     5472   167927   +162455     
===========================================
+ Hits         4520   135010   +130490     
- Misses        759    28072    +27313     
- Partials      193     4845     +4652     

Flag	Coverage Δ
cli-ort-linux	82.60% <ø> (ø)
cli-ort-windows	82.19% <ø> (?)
offline	80.32% <96.68%> (?)

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

Files with missing lines	Coverage Δ
...nnx-runtime-ep-cpu/src/kernels/simd_activations.rs	90.98% <96.68%> (ø)

... and 363 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.