perf(cpu): route RoPE and Softmax through the CPU task runtime by justinchuby · Pull Request #1202 · justinchuby/onnx-genai

justinchuby · 2026-08-18T04:38:44Z

What

Routes RotaryEmbedding and Softmax off Rayon and onto the CPU task runtime (#1201). Both previously sized their fan-out from rayon::current_num_threads() and split with par_chunks_mut, paying Rayon's park latency on every region and starting a second pool beside ORT's inside a plugin-EP call. They now go through task_runtime::chunks_mut, which lands on ORT's intra-op pool inside a compute call and on the native adaptive-spin pool outside. Grain policy changes from a static per-worker share to a dynamic-claim floor.

Review notes (validated on Intel i7-13800H, 14 physical / 20 logical, Windows; baseline = origin/main)

Correctness / bit-identity. This is a pure-scheduling change (reordering parallel work), so output is unchanged. Verified by the existing bitwise kernel tests, which all pass: rotary_*_bit_identical, rotary_*_bitwise, rotary_strided_cache_fallback_matches_contiguous_fast_path_bitwise, and the softmax/log-softmax reference tests. The range-partition hazard is guarded by task_chunking_covers_every_unit_exactly_once, which drives the real partition through task_runtime::chunks_mut and asserts each unit is written exactly once.
Reachability. RoPE and Softmax now call task_runtime::chunks_mut/chunk_runs_mut; rayon/par_chunks is fully removed from both files — not an unwired placeholder.
Cross-cluster conflict (activation kernels). perf(cpu): route RoPE and Softmax through the CPU task runtime #1202's original commit 773d6a94 force-inlined map_ps/map_bias_ps with #[inline(always)] (no target_feature) to stop the RoPE/Softmax rewiring de-vectorizing tanh_avx2/sigmoid_avx2 via a whole-crate inline-cost-model flip. The activation cluster has since landed the competing fix on main (#[inline] + #[target_feature(avx2,fma)], mutually exclusive with inline(always)). Per the no-revert directive I kept main's version and dropped 773d6a94, then verified by disassembly that the de-vectorization does not recur: with perf(cpu): route RoPE and Softmax through the CPU task runtime #1202's routing applied, tanh_avx2 = 124 lines / 0 calls / 68 ymm ops and sigmoid_avx2 = 127 / 0 / 69 — both fully vectorized, relu/erf inlined away. So the two clusters are compatible and perf(cpu): route RoPE and Softmax through the CPU task runtime #1202 no longer touches simd_activations.rs.
Tests: cargo test -p onnx-runtime-ep-cpu --lib = 1405 passed / 0 failed / 17 ignored. Clippy --all-targets -D warnings clean.
Perf. The park-latency win is inherited from perf(cpu): add a CPU task runtime with an adaptive-spin native pool #1201's reproduced dispatch-latency result (p50 1.9 µs flat across 0–100 µs gaps on this box). The PR's per-shape RoPE/Softmax ORT ratios were measured on an AMD EPYC 9V74 and are not re-run here — they need the scripts/ort_ab ORT session harness and are hardware-dependent by nature.

Rebased onto latest main; benchmark writeup renumbered to fit after the erf section.

github-actions · 2026-08-18T13:43:39Z

🔴 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_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1`	63.43 µs	125.21 µs	+97.4%
🔴	`gather/medium_f32_threads=1-internal/32768`	4.23 µs	7.76 µs	+83.7%
🔴	`add/small_f32_threads=1-internal/1024`	197.3 ns	322.9 ns	+63.7%
🔴	`block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024`	78.94 µs	124.98 µs	+58.3%
🔴	`add/small_f16_threads=1-internal/1024`	475.1 ns	711.6 ns	+49.8%
🔴	`block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1`	419.28 µs	619.14 µs	+47.7%
🔴	`matmul/small_generic_f32_threads=1/1x256x256`	35.85 µs	51.39 µs	+43.4%
🔴	`matmul/large_generic_f32_threads=8/32x1024x1024`	4.78 ms	6.53 ms	+36.6%
🔴	`matmul/medium_generic_f32_threads=8/32x512x512`	1.46 ms	1.93 ms	+32.6%
⚠️	`matmul/large_generic_f16_threads=1/32x1024x1024`	88.13 µs	112.20 µs	+27.3%
⚠️	`matmul/small_generic_f32_threads=8/1x256x256`	42.50 µs	53.29 µs	+25.4%
⚠️	`matmul/small_generic_f16_threads=8/1x256x256`	30.93 µs	38.25 µs	+23.7%
⚠️	`matmul/medium_generic_f16_threads=1/32x512x512`	33.12 µs	40.25 µs	+21.5%
⚠️	`add/medium_f16_threads=1-internal/262144`	115.16 µs	134.02 µs	+16.4%
✅	`matmul/medium_generic_f16_threads=8/32x512x512`	39.28 µs	44.82 µs	+14.1%
✅	`sampling_latency/top_p_per_token`	422.77 µs	479.41 µs	+13.4%
✅	`add/medium_f32_threads=1-internal/262144`	28.30 µs	31.93 µs	+12.8%
✅	`add/large_bf16_threads=1-internal/4194304`	2.25 ms	2.53 ms	+12.2%
✅	`reduce_mean/small_f32_threads=1-internal/4096`	14.77 µs	16.50 µs	+11.7%
✅	`tokenization/encode_tokens_per_second`	459.95 µs	512.62 µs	+11.5%
✅	`add/medium_bf16_threads=1-internal/262144`	103.21 µs	114.39 µs	+10.8%
✅	`gather/large_bf16_threads=1-internal/131072`	13.97 µs	15.43 µs	+10.4%
✅	`gather/large_f16_threads=1-internal/131072`	15.16 µs	16.68 µs	+10.0%
✅	`matmul/small_generic_f16_threads=1/1x256x256`	32.63 µs	35.37 µs	+8.4%
✅	`matmul/large_generic_bf16_threads=1/32x1024x1024`	2.10 ms	2.26 ms	+7.6%
✅	`block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024`	1.10 ms	1.18 ms	+7.2%
✅	`matmul/medium_generic_bf16_threads=8/32x512x512`	581.01 µs	623.05 µs	+7.2%
✅	`add/large_f32_threads=1-internal/4194304`	740.34 µs	779.30 µs	+5.3%
✅	`add/small_bf16_threads=1-internal/1024`	471.4 ns	493.6 ns	+4.7%
✅	`gather/small_bf16_threads=1-internal/4096`	518.7 ns	542.8 ns	+4.6%
✅	`gather/small_f16_threads=1-internal/4096`	551.8 ns	571.3 ns	+3.5%
✅	`matmul/large_generic_f16_threads=8/32x1024x1024`	98.20 µs	99.23 µs	+1.1%
✅	`gather/small_f32_threads=1-internal/4096`	729.1 ns	736.6 ns	+1.0%
✅	`matmul/medium_generic_bf16_threads=1/32x512x512`	591.32 µs	596.96 µs	+1.0%
✅	`sampling_latency/greedy_per_token`	3.44 µs	3.47 µs	+0.9%
✅	`matmul/large_generic_bf16_threads=8/32x1024x1024`	2.56 ms	2.54 ms	-0.7%
✅	`qwen3_sampling_processors/top_p_fast_after_top_k`	617.14 µs	611.61 µs	-0.9%
✅	`block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024`	99.17 µs	97.86 µs	-1.3%
✅	`reduce_mean/large_f32_threads=1-internal/262144`	1.06 ms	1.05 ms	-1.4%
✅	`reduce_mean/medium_f32_threads=1-internal/65536`	265.20 µs	260.37 µs	-1.8%
✅	`qwen3_sampling_processors/top_k_full_sort_baseline`	2.32 ms	2.23 ms	-3.9%
✅	`matmul/large_generic_f32_threads=1/32x1024x1024`	11.03 ms	10.59 ms	-4.0%
✅	`sampling_latency/top_k_per_token`	58.33 µs	55.90 µs	-4.2%
✅	`tokenization/decode_tokens_per_second`	8.47 ms	8.09 ms	-4.5%
✅	`logit_processing/seven_processor_chain_per_step`	391.05 µs	364.97 µs	-6.7%
✅	`gather/medium_f16_threads=1-internal/32768`	3.01 µs	2.79 µs	-7.3%
✅	`matmul/small_generic_bf16_threads=1/1x256x256`	38.46 µs	34.75 µs	-9.6%
✅	`add/large_f16_threads=1-internal/4194304`	2.15 ms	1.94 ms	-10.0%
✅	`matmul/small_generic_bf16_threads=8/1x256x256`	41.72 µs	37.27 µs	-10.7%
✅	`qwen3_sampling_processors/top_k_partial_selection`	216.90 µs	192.29 µs	-11.3%
✅	`qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline`	4.68 ms	4.11 ms	-12.1%
🟢	`qwen3_sampling_processors/top_k_top_p_full_sort_baseline`	7.17 ms	5.99 ms	-16.4%
🟢	`sampling_latency/min_p_per_token`	285.75 µs	235.24 µs	-17.7%
🟢	`matmul/medium_generic_f32_threads=1/32x512x512`	3.07 ms	2.51 ms	-18.1%
🟢	`qwen3_sampling_processors/top_k_top_p_fast`	844.00 µs	675.07 µs	-20.0%
🟢	`kv_cache/alloc_dealloc_pages`	50.97 µs	40.46 µs	-20.6%
🟢	`gather/medium_bf16_threads=1-internal/32768`	3.27 µs	2.58 µs	-21.1%
🟢	`grammar_masking/llguidance_compute_mask/32`	96.95 µs	71.56 µs	-26.2%
🟢	`gather/large_f32_threads=1-internal/131072`	47.93 µs	35.01 µs	-26.9%

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.44 4.15 5.10 }

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)

## What `main` at 4d231ea fails two blocking CI steps on current stable (rustc/rustfmt 1.9.0, 1.97.1): ``` $ cargo fmt --all -- --check Diff in crates/onnx-runtime-ep-cuda/src/provider.rs Diff in crates/onnx-runtime-memory-api/src/allocator.rs Diff in crates/onnx-runtime-memory-api/src/capability.rs $ cargo clippy --locked --all-targets -p onnx-runtime-session -- -D warnings error: field `0` is never read --> crates/onnx-runtime-session/src/executor/mod.rs:175:45 ``` Both landed while the runner queue was saturated (71 queued / 1 in progress at the time of writing, nothing completed on `main` since 08:18Z), so no PR has seen a red check yet. Every open PR in the repo currently inherits both failures. ## Why these fixes **rustfmt** — mechanical normalisation, no semantic change. **`ActivationPlanForTest`** (from #1226) is a tuple struct whose single field is the `globals_lock()` `MutexGuard`. `dead_code` does not model "this field's value is its `Drop`", so it fires. The guard must stay: releasing it early is exactly the leaked-planner-gate race the struct was added to prevent. So the lint is silenced with a comment explaining the RAII intent, rather than the field removed. ## Verification | check | result | | --- | --- | | `cargo fmt --all -- --check` | clean | | `cargo clippy --locked --all-targets -p onnx-runtime-session -- -D warnings` | clean | | `cargo test --locked -p onnx-runtime-session --lib` | 181 passed, 0 failed | Found while trying to land the CPU task-runtime stack (#1201 → #1202 → #1207 → #1232 → #1238); this is unrelated to that work and is deliberately kept out of it so it can go in on its own. Working as sebastian (Performance Engineer) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

…1201) ## What Adds a CPU task runtime (`task_runtime`) with an adaptive-spin native pool, plus an SMT/hybrid-aware `core_topology` module. This is the *foundation* PR: it introduces the machinery but wires no production kernel to it yet (RoPE/Softmax routing is #1202). `matmul_nbits` only gains a `decode_thread_budget()` getter. ## Design - **One pool per machine, not per subsystem.** Inside an ORT plugin-EP compute call, fan-outs route to `KernelContext_ParallelFor` (`Backend::Host`) and the native pool stays parked, so the new pool never spins alongside ORT's intra-op pool. Outside ORT it uses the native adaptive-spin pool (`Backend::Native`); tiny work runs `Backend::Serial`. - **Width is capped to the physical cores the process may run on** (min floor 8), and reuses the decode thread budget so one knob sizes the engine. - **Adaptive spin** grows the spin window while dispatches keep arriving and decays to a park when they stop. ## Review notes (validated on this hardware) **Hardware: Intel i7-13800H, 14 physical / 20 logical (6 P-cores×2 SMT + 8 E-cores), Windows.** Baseline reference = origin/main. - **Correctness / thread count.** In this PR no production kernel dispatches to the pool (rayon remains everywhere), so a running inference gains **zero** steady-state threads from it; the pool is lazily built on first use only. Verified there are no `task_runtime::` callers outside the module + tests. - **Bug found & fixed during review.** The Windows `GroupMask` offset in `core_topology` was computed as the address delta between two unrelated stack copies, so `GetLogicalProcessorInformationEx` detection returned `None` on this box. The SMT/physical-core cap then silently no-op'd and the pool resolved to **20 workers on 14 physical cores** — the exact oversubscription this exercise exists to remove. Fixed with `offset_of!`; detection now reports `core_count=14`, hybrid split `[2×6, 1×8]`, and the resolved width drops **20 → 14**. Added a non-vacuous Windows regression test (the pre-existing self-consistency test returned early on `None` and was vacuous on Windows). - **Dispatch latency reproduced** (release, 14 workers, 400 rounds), `tests/task_runtime_latency.rs`: | gap before dispatch | p50 | p90 | | concurrent sessions | p50 | p90 | |---|---|---|---|---|---|---| | 0 µs | 1.7 µs | 1.9 µs | | 1 | 2.0 µs | 2.2 µs | | 20 µs | 1.9 µs | 2.1 µs | | 2 | 1.9 µs | 2.4 µs | | 100 µs | 1.9 µs | 2.6 µs | | 4 | 2.3 µs | 2.7 µs | | 500 µs | 142.9 µs | 186.5 µs | | 8 | 2.5 µs | 3.2 µs | The point is the flatness: 0 µs and 100 µs cost the same, so a decode step no longer pays for having been idle; past the 500 µs spin ceiling the workers park and the cost returns to a futex wake, which is correct. `slot exhaustion: 0 of 9744 dispatches fell back to serial`. (These are on this Intel laptop; the PR's original 4.8 µs figures were on an AMD EPYC 9V74 server — peers, not a contradiction.) - **Tests:** `cargo test -p onnx-runtime-ep-cpu --lib` = 1395 passed / 0 failed (baseline main = 1353; +42 new). Task-runtime integration tests pass individually. Clippy `--all-targets -D warnings` clean. Note: the idle test skips on Windows (no per-process CPU accounting), and one lib timing test flaked once when the whole suite ran concurrently with the idle-window binary under full CPU load — it passed on every isolated and repeat run; correctness never depends on the pool (serial fallback). Merged onto latest main (resolved a benchmark-doc section-numbering conflict; the task-runtime writeup is §33). --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Both kernels sized their fan-out from `rayon::current_num_threads()` and split with `par_chunks_mut`, so both paid Rayon's park latency on every region -- 67us back-to-back vs 226us after a 20us gap, which is the shape of a decode step -- and both ignored the plugin EP's borrowed ORT pool entirely. Route them through `task_runtime::chunks_mut` instead. Inside a plugin-EP compute call the fan-out now lands on ORT's intra-op pool rather than starting a second one beside it; outside, it lands on the native adaptive-spin pool. The grain policy changes shape as well as owner. Both kernels used to compute a *task size* from the pool width -- RoPE aimed at four tasks per worker, Softmax at exactly one -- which assigns each worker a static share and makes the whole region wait for whichever share was unlucky: a descheduled thread, an SMT sibling, a cold row. The runtime claims tasks dynamically, so the kernels now state only a *floor* (the smallest run worth handing to another thread) and let the runtime pick the count. RotaryEmbedding: `rotary_units_per_task` becomes `rotary_min_units_per_task`, which keeps the measured MIN_ROTARY_TASK_ELEMENTS bar (16 Ki elements; chunking per layout unit instead measured slower than serial at every thread count) and drops the width term. Softmax: the chunk becomes one ROW_TILE_BYTES row tile instead of `n / workers`. That is the size the serial worker already tiles at, to keep a run of rows in a core's private cache between the scale/mask write and the softmax's two reads, so a chunk smaller than a tile could not fill that pipeline and would call MLAS more often for less work each time. `task_chunking_covers_every_unit_exactly_once` now drives the real partition through `task_runtime::chunks_mut` and checks each unit was written by exactly one task, rather than checking arithmetic that no longer decides the tiling. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

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

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

justinchuby · 2026-08-18T16:16:18Z

Reviewed and validated on Intel i7-13800H (14 physical / 20 logical), Windows; baseline = origin/main.

Pure-scheduling change (rayon -> task_runtime for RoPE/Softmax). Output bit-identity preserved: all bitwise rotary + softmax/log-softmax reference tests pass; range partitioning guarded by task_chunking_covers_every_unit_exactly_once (drives the real task_runtime::chunks_mut partition, asserts each unit written once). Routing is reachable -- rayon/par_chunks fully removed from both files.

Cross-cluster conflict in simd_activations.rs: #1202's original 773d6a9 force-inlined map_ps/map_bias_ps with #[inline(always)] (no target_feature) to prevent the RoPE/Softmax rewiring de-vectorizing tanh/sigmoid. The activation cluster landed the competing #[inline]+#[target_feature(avx2,fma)] fix on main (mutually exclusive with inline(always)). Per the no-revert directive I kept main's version and dropped 773d6a9, then verified by disassembly that de-vectorization does NOT recur with #1202's routing applied: tanh_avx2 = 124 lines/0 calls/68 ymm, sigmoid_avx2 = 127/0/69 -- both fully vectorized. #1202 therefore no longer touches simd_activations.rs; the two clusters are compatible.

Tests: lib 1405 passed/0 failed; clippy --all-targets -D warnings clean. Perf park-latency win inherited from #1201's reproduced dispatch result; per-shape ORT ratios were AMD-EPYC and need the ort_ab harness, not re-run here.

Verdict: GO.

codecov · 2026-08-18T17:36:40Z

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 80.04%. Comparing base (774b256) to head (28c75ea).
⚠️ Report is 24 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1202      +/-   ##
==========================================
+ Coverage   79.69%   80.04%   +0.35%     
==========================================
  Files         358      363       +5     
  Lines      155271   159952    +4681     
  Branches   155271   159952    +4681     
==========================================
+ Hits       123743   128037    +4294     
- Misses      26943    27266     +323     
- Partials     4585     4649      +64

Flag	Coverage Δ
mlas	`85.09% <ø> (?)`
offline	`79.94% <100.00%> (+0.25%)`	⬆️

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/rotary_embedding.rs	`93.63% <100.00%> (-0.01%)`	⬇️
crates/onnx-runtime-ep-cpu/src/kernels/softmax.rs	`92.33% <100.00%> (-0.05%)`	⬇️

... and 22 files with indirect coverage changes

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

…1232) ## The CPU budget was buying hyperthreads, not cores `ONNX_GENAI_CPU_DECODE_THREADS=N` confines the process to `N` logical CPUs. `choose_budget_cpus` picked them as the `N` lowest indices of the chosen NUMA node. Every host we run on numbers SMT siblings adjacently — `0-1`, `2-3`, … on AMD EPYC and on Intel since Skylake-SP — so **a budget of `N` landed on `N/2` physical cores**. The symptom was unmistakable once measured. Int4 `MatMulNBits` 4096×6144, 128 tokens, native-only, 20 runs, on the 16-core/32-thread EPYC 9V74: | budget | wall per run | user CPU | | ------ | ------------ | -------- | | 1 | 79.4 ms | 2.04 s | | 2 | 81.2 ms | **4.10 s** | | 3 | 40.4 ms | 2.05 s | | 4 | 53.9 ms | 3.35 s | A budget of 2 was exactly as slow as a budget of 1 while burning two cores. A budget of 3 beat a budget of 4. `taskset` confirms the cause directly — `0,2` (two cores) ran in 50.1 ms against `0,1` (one core, two threads) at 81.8 ms, using 36% less CPU. ## The fix A **ranking, not a widening**: * `scatter_across_cores` orders the candidate pool by `CoreTopology::leaders_within` — one CPU per physical core first, siblings after — and truncates *after* ranking. The result stays a subset of the same pool, so the cpuset/cgroup guarantee is unchanged and a full-width budget returns exactly what it did before. * `smt_scaled_request` sizes the NUMA node search in cores rather than logical CPUs, falling back to the old logical sizing when no node is that large, so a budget that fits on one core-rich node still stays on one node. * the cross-node top-up now happens *before* ranking, so a second node's fresh cores can displace the first node's SMT siblings instead of being appended to an already-full mask. * `order_pin_targets` applies the same order to the flat decode Rayon pool's pin list (the builder pins worker `i` to `cpus[i % len]`, so the order decides whether an 8-worker pool occupies 8 cores or 4). The persistent SPMD pool and the `numa-split` sub-pools are deliberately left compact — their workers spin, and spreading spinning workers one per core is the experiment already recorded in `core_topology`'s module docs, where it measured *worse*. `choose_budget_cpus` gained a `cores: Option<&CoreTopology>` parameter so the policy stays pure and unit-testable; with `None` (no discoverable SMT map) every function here is the identity and the old behaviour stands exactly. ## Result 42 cells (20 GEMM, 22 transform) × 6 widths × 5 trials, paired arms in one driver invocation. Native-only view — the mask moves ORT's threads too, so only native-vs-native isolates the change. Geomean of `before/after`, >1 is faster: | width | GEMM | >1.05× | <0.95× | Transforms | >1.05× | <0.95× | | ----- | ---- | ------ | ------ | ---------- | ------ | ------ | | 1 | 0.996 | 0 | 1 | 0.997 | 1 | 3 | | 2 | **1.769** | 19 | **0** | 1.191 | 12 | **0** | | 4 | **1.642** | 20 | **0** | 1.095 | 11 | 2 | | 8 | **1.441** | 16 | 1 | 0.913 | 4 | 9 | | 16 | **1.245** | 15 | 1 | 0.961 | 6 | 9 | | 32 | 0.976 | 8 | 6 | 1.062 | 9 | 4 | The two ends are the control: a budget of 1 cannot be re-ranked and a budget of 32 is the whole machine, so both must be flat — and both are. At 2 and 4 threads **not one cell out of 42 regressed**. The transform column at t=8/16 reads as a loss and is not one: those cells are 20–500 µs at five trials on a shared box. Re-measured at 40 runs × 5 warmups with the arms strictly interleaved, they are wins — `sm_bert_b8_s128` by 1.9× at both widths, `rope_gptj_il_s512` by 1.35× at t=8. The five-trial grid is left in the docs unedited so the distinction stays visible; §34.5 has the raw pairs. ## Validation * `cargo test -p onnx-runtime-ep-cpu` — 1391 passed, 0 failed * `cargo clippy -p onnx-runtime-ep-cpu --all-targets` — clean * `cargo fmt --all --check` — clean * 8 new unit tests, all of which fail if the ranking is reverted, including the cpuset-containment case (a leader that exists in the topology but not in the process's allowed set must not be invented), the asymmetric-node case that exercises `smt_scaled_request` through `choose_budget_cpus`, the top-up-before-ranking case, and the "a budget that fits one node does not spill across nodes for cores" case. Documented as **§34 Phase 14** in `docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md`. ## Rebase and convergence (2026-08-18) #1201/#1202/#1207 were **squash-merged**, so this branch's ancestry no longer reached `main`. Its three commits were cherry-picked onto `main` at `c55a3fab3`; all applied cleanly, and the diff is now self-contained: 3 files, +465/−29. Review the whole PR, not "the last two commits". Defect found and fixed during the rebase: the benchmark section was numbered `## 32. Phase 12`, which **collided with the existing §32 (Phase 12, erf Estrin)** on `main`. Renumbered to `## 34. Phase 14` (subsections 34.1–34.6), and the internal back-reference "Phase 11 built a task runtime" corrected to "Phase 13 (§33) built a task runtime" — §33 is where the task runtime actually landed. The squad decision note's `§32` cross-reference was updated to match. Re-validated on the rebased head (rustc 1.97.1, the same toolchain CI resolves): * `cargo test -p onnx-runtime-ep-cpu --lib` — **1426 passed, 0 failed, 17 ignored** * `cargo clippy -p onnx-runtime-ep-cpu --all-targets` — clean * `cargo fmt --all -- --check` — clean CI note: the repository's Actions queue is saturated (every recent run is `queued`, nothing `in_progress`), so the two required checks — `Fast (Linux x86_64)` and `Rust quality` — cannot report. Both lanes were reproduced locally, step for step, from `.github/workflows/ci.yml`. `main` itself fails **both** of them today; #1346 is the fix for that and lands first. --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

perf(cpu): run the MLAS prefill tiling on the CPU task runtime The `m > 1` MLAS SQNBit prefill tiling (`run_mlas_shards`) issued a single `tiles.par_iter()` fan-out on global Rayon, sized by `rayon::current_num_threads()`. With a co-resident ORT intra-op pool spinning on the same cores, every parked-Rayon wake-up lands behind a spinning thread, which is what made `gemm_nbits_*_t8` at t=32 the worst cell in the benchmark ledger. This routes that fan-out through the CPU task runtime from #1201 instead, with a work-size policy for the one case where the SMT-capped pool leaves hardware threads idle. ## MLAS remains opt-in and non-load-bearing This does **not** enable MLAS anywhere. The changed code lives entirely inside the pre-existing `m > 1 && active > 1 && !mlas_prefill_serial()` path that already called `mlas_sys::sqnbit_gemm_into`. No Cargo feature, `#[cfg(feature = "mlas")]` gate, or default-feature set is touched: `mlas` is still opt-in (`default = ["full"]`, and `full` does not include `mlas`). The new policy helpers (`prefill_fan_out`, `prefill_tile_grain`, `PrefillFanOut`) are pure integer arithmetic marked `#[cfg_attr(not(feature = "mlas"), allow(dead_code))]` so they compile and their unit tests run on the default (mlas-off) CI lane even though only the MLAS path consults them. MLAS stays a labelled reference arm. ## What changed - `prefill_fan_out(macs, lanes, wide)`: below `WIDE_PREFILL_MACS` (512 Mi MACs), or whenever global Rayon is not actually wider than the pool, fan out on the task runtime (cheap ~5 us dispatch, topology-aware, no fight with a co-resident ORT pool). Above it, use the wider global Rayon path -- a prefill tile is a multi-ms MLAS call whose dequantise step has enough load latency that SMT siblings pay off, so the SMT cap costs more than a 226 us park wake-up (0.25% of a 90 ms fan-out). - `prefill_tile_grain`: a per-task tile floor so no task gets less than `MIN_PREFILL_TASK_MACS` (512 Ki) of arithmetic. ## Verification Hardware: Intel Core i7-13800H, 14 physical / 20 logical (6 P + 8 E). Baseline: `main` at the rebase point. Toolchain: cargo 1.97.1. - **Bit-identity (the important one).** This is pure scheduling: the `run_tile` closure and the `sqnbit_gemm_into` call are byte-for-byte the same, only the executor and grain differ, and every tile writes a disjoint `[row, row+rows) x [shard.start, shard.start+len)` window so order cannot matter. Confirmed empirically under `--features mlas`: `mlas_prefill_parallel_dispatch_matches_serial` and `mlas_prefill_dispatch_parity_subprocess` pass -- the routed parallel tiling matches the serial reference. - **Policy tests (default features).** All six `prefill_*` unit tests pass. - **Falsified.** Flipping the threshold comparison in `prefill_fan_out` from `<` to `<=` turns `large_prefill_work_takes_the_wide_fan_out` RED (`left: TaskRuntime, right: Wide` at exactly `WIDE_PREFILL_MACS`) -- the test is non-vacuous and guards the boundary. Restored to green. - **`--features mlas` compiles clean;** clippy `--all-targets -D warnings` clean on default features. ## Perf The §35 (Phase 15) tables in the benchmark doc record up to 13.5x at t=32 on the small int4 cells, dropping `gemm_nbits_*_t8` from 22-34x ORT to 2.6-2.8x. Those tables mix the author's EPYC 9V74 (16c/32t) and the laptop measurements; per the repo's measurement rule they are peers, named by hardware. The mechanism (the work-size policy, the grain floor, the disjoint-window safety) and correctness are verified, and the policy decisions are reproduced in unit tests; no full ORT A/B sweep was re-run on the laptop, so the headline speedup magnitudes are the author's, not independently re-measured there. ## Rebase and convergence (2026-08-18) #1201/#1202/#1207 and #1143 were **squash-merged**, so this branch's ancestry no longer reached `main`. Its three commits were cherry-picked onto `main` at `c55a3fab3` and applied cleanly. The diff is now self-contained: 2 files, +355/-25 (`matmul_nbits.rs` and the benchmark doc). It is **not** stacked on #1232 any more. Section numbering: #1232 lands first and takes §34/Phase 14, so this PR's section was renumbered to **§35/Phase 15**. (The `34.1×` figure in the §35.4 matrix is a speed ratio, not a section reference, and is unchanged.) Re-validated on the rebased head (rustc 1.97.1, the toolchain CI resolves): * `cargo test -p onnx-runtime-ep-cpu --lib` — **1433 passed, 0 failed, 17 ignored** (on top of #1232) * `cargo clippy -p onnx-runtime-ep-cpu --all-targets` — clean * `cargo fmt --all -- --check` — clean * under `--features mlas`, the parity test that actually guards this change, `mlas_prefill_parallel_dispatch_matches_serial`, **passes** ### The two `--features mlas` failures are pre-existing on `main` Running with `--features mlas` fails two tests: `feature_default_guard::mlas_is_not_a_default_feature` (it panics *because* `--features mlas` was passed explicitly) and `kernels::simd_activations::mlas_ab::mlas_matches_rust_simd_on_special_values` (a 1-ULP Erf disagreement between the MLAS and pure-Rust SIMD routes). **Both reproduce identically on `main` at `c55a3fab3` with this PR's changes absent**, so they are baseline-equivalent, not regressions, and they are out of scope here. Neither runs on any required lane: `mlas` is not a default feature. ### The red-criterion comment is runner noise The criterion report flags regressions including `tokenization/decode_tokens_per_second` -47.6%. This diff cannot reach the tokenizer, and every changed line of `matmul_nbits.rs` is inside the `#[cfg(feature = "mlas")]` `run_mlas_shards` path, which the benchmark build (default features) does not compile in. The default-feature build is behaviourally identical to `main`; the deltas are shared-runner variance. ### CI The repository's Actions queue is saturated (every recent run is `queued`, nothing `in_progress`), so the two required checks — `Fast (Linux x86_64)` and `Rust quality` — cannot report. Both lanes were reproduced locally, step for step, from `.github/workflows/ci.yml`. `main` itself fails **both** of them today; #1346 is the fix for that and lands first. Closes #1238. Working as sebastian (CPU perf). --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

justinchuby force-pushed the squad/sebastian-cpu-task-runtime-kernels branch 2 times, most recently from c1390f8 to 3cf4f18 Compare August 18, 2026 04:57

This was referenced Aug 18, 2026

perf(cpu): route Transpose and the elementwise fallback through the task runtime #1207

Merged

perf(cpu): add a CPU task runtime with an adaptive-spin native pool #1201

Merged

justinchuby force-pushed the squad/sebastian-cpu-task-runtime-kernels branch from 3cf4f18 to e1ca38a Compare August 18, 2026 05:31

justinchuby force-pushed the squad/sebastian-cpu-task-runtime branch from 8d18014 to 6b14779 Compare August 18, 2026 05:34

justinchuby force-pushed the squad/sebastian-cpu-task-runtime-kernels branch 9 times, most recently from a4395c9 to bab6433 Compare August 18, 2026 07:13

justinchuby marked this pull request as ready for review August 18, 2026 07:16

justinchuby force-pushed the squad/sebastian-cpu-task-runtime-kernels branch from bab6433 to 8fe0695 Compare August 18, 2026 07:35

justinchuby mentioned this pull request Aug 18, 2026

perf(cpu): spread the decode CPU budget one thread per physical core #1232

Merged

justinchuby force-pushed the squad/sebastian-cpu-task-runtime-kernels branch from 26a673c to cd6e3e1 Compare August 18, 2026 09:12

This was referenced Aug 18, 2026

fix(ci): unbreak the Rust quality lane on current stable #1239

Merged

perf(cpu-ep): zero-copy activation fast path and an AVX2 Elu kernel #1240

Merged

Base automatically changed from squad/sebastian-cpu-task-runtime to main August 18, 2026 16:00

justinchuby and others added 3 commits August 18, 2026 09:03

docs(bench): record the RoPE and Softmax matrices, and the de-vectori…

0e5bfa5

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

docs(bench): re-verify the Softmax win against the merged #1219 baseline

28c75ea

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

justinchuby force-pushed the squad/sebastian-cpu-task-runtime-kernels branch from cd6e3e1 to 28c75ea Compare August 18, 2026 16:14

justinchuby merged commit 508de85 into main Aug 18, 2026
7 of 17 checks passed

justinchuby deleted the squad/sebastian-cpu-task-runtime-kernels branch August 18, 2026 16:16

justinchuby mentioned this pull request Aug 19, 2026

perf(cpu): run the MLAS prefill tiling on the CPU task runtime #1238

Merged

Navigation Menu

Repository navigation

perf(cpu): route RoPE and Softmax through the CPU task runtime - #1202

perf(cpu): route RoPE and Softmax through the CPU task runtime#1202
justinchuby merged 3 commits into
mainfrom
squad/sebastian-cpu-task-runtime-kernels

justinchuby commented Aug 18, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Aug 18, 2026 •

edited

Loading

Uh oh!

justinchuby commented Aug 18, 2026

Uh oh!

Uh oh!

codecov Bot commented Aug 18, 2026 •

edited

Loading

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!

What

Review notes (validated on Intel i7-13800H, 14 physical / 20 logical, Windows; baseline = origin/main)

Uh oh!

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

Uh oh!

🔴 Benchmark Regression Detected

Uh oh!

justinchuby commented Aug 18, 2026

Uh oh!

Uh oh!

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

Uh oh!

Codecov Report

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

justinchuby commented Aug 18, 2026 •

edited

Loading

github-actions Bot commented Aug 18, 2026 •

edited

Loading

codecov Bot commented Aug 18, 2026 •

edited

Loading