Repository navigation
perf(cpu-ep): vectorise RoPE and size its Rayon tasks by work, not by layout - #1175
Merged
Merged
Conversation
… layout The RotaryEmbedding kernel fanned out one task per layout unit - one `[B,H,S,D]` plane or one `[B,S,H*D]` row. On `rope_llama3_s128` that is 4096 elements (16 KiB) per task, and the scheduling cost swamped it: the parallel path was *slower* than the serial one at every thread count (t=1 162us vs 126us, t=8 142us vs 126us, t=16 195us vs 127us). The kernel also left its inner loop to the autovectoriser, which did not take it. Two changes: * `rotary_units_per_task` aggregates whole layout units until each task clears 16 Ki elements, targeting ~4 tasks per worker for balance, and collapses to a single serial pass when the pool has one worker. * `rotate_split_half` / `rotate_interleaved` gain explicit AVX2 arms. They use separate multiply/add rather than FMA so the vector body and the scalar tail round identically - a tensor must not change value because a different thread count split it in a different place. Measured `exec_kernel.compute` p50 on `rope_llama3_s128`, native-only: 242 -> 85us (t=1), 174 -> 67us (t=8), 138 -> 59us (t=16). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The transform grid only covered `interleaved=0`. The GPT-J convention rotates adjacent even/odd channels, which is a structurally different inner loop, so it needs its own cells in the A/B matrix. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…ound Opus adversarial review of #1175 verified the permutation algebra, the target_feature gating and the chunking arithmetic, and raised four non-blocking findings. All four applied: * The SAFETY comments claimed an "AVX2+FMA probe"; the probe checks `avx2` only, and deliberately so - the arms use no FMA intrinsic. * The `half` sweeps only reached residues {0,1,3,7} mod 8, so the comment's "every residue" claim was false. Both sweeps now cover 1..9 plus 13, 14, 20 so every residue is hit before and after several full vector iterations. Falsified: deleting either scalar tail now fails at half=1 (it did not before). * No shape drove the parallel re-chunk with `total_units` indivisible by `per_task`, so the short final slab was only proved correct by hand. Added (1,130,32,128) 3D and (2,66,17,64) 4D. Falsified: an off-by-one task stride now fails the sweep. * `vector_rotation_available` is used only in this module; narrowed from `pub(crate)` to private. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
marked this pull request as ready for review
August 18, 2026 02:48
justinchuby
added a commit
that referenced
this pull request
Aug 18, 2026
) ## Why Two problems with the master benchmark doc as it stands on `main`: 1. **§21.2/§21.3 are stale.** They describe the input-binding copy as an open gap and explicitly decline to say what closing it would be worth. #1146 closed it — `bind_inputs` on `rope_llama3_s128` went from ~601 µs to 0.7 µs. Every ratio in the document that predates that merge is being read against the wrong baseline, and nothing in the document says so. 2. **§23.7 frames the remaining losses as five independent kernel problems.** They are mostly one problem, and it is not in any of those kernels. ## What this adds ### §24 — Phase 6 (#1146) What the borrowed `DeviceBuffer` does and the exact preconditions; the sequence-storage choke point (`read_seq_element` *moves* a buffer into cross-run `shared_buffers`, so guarding the slot was not enough); the 57-cell result (56/57 improved, 342/342 parity PASS); and precisely which §21 claims it supersedes. ### §25 — Phase 7 (#1175) The two independent `RotaryEmbedding` defects, with the measurement that made them visible once #1146 left `exec_kernel.compute` at 96% of the run: * the fan-out was **net-negative at every thread count** (t=16: 195 µs parallel vs 127 µs serial; 36 µs of pure Rayon bookkeeping to fan 128 tasks over a *one-worker* pool); * the inner loop was scalar at ~1.4 cycles/element. Plus why the new AVX2 arms deliberately avoid FMA (a tensor must not change value because a different thread count split it in a different place), and the 21-cell matrix: 21/21 improved, 0 regressed, 5 now beat ORT. ### §26 — the root cause behind the `t=8/16` column An isolated 524 288-element region on 16 workers, varying only the idle gap between regions: | gap | region cost | |---|---| | none | 67 µs | | 20 µs | 226 µs | | 500 µs | 236 µs | A 20 µs gap — less than the serial glue between two graph nodes — costs ~160 µs. Rayon parks its workers; ORT's intra-op pool spins first, so its workers are hot when the next op arrives and ours are not. `rope_llama3_s128` at t=16 is 59 µs of kernel against ORT's ~30 µs, but a 189 µs paired wall. `decode_spmd` already implements the right policy (bounded `KMP_BLOCKTIME` spin) and is wired only to the decode GEMM path. A prototype routing RoPE through it takes t=16 from **79 µs to 25 µs native-only, past ORT's ~30 µs**. It is documented rather than merged, for two reasons that are both worth recording: **(a) A latent correctness hazard.** `SharedState::publish` writes the job into a single `UnsafeCell` slot (`decode_spmd.rs:243`); `wait`'s doc comment states "the dispatcher is a single, never-idle thread" as an assumption, not an invariant. Two concurrent dispatchers overwrite each other's job pointer and both sets of workers run whichever closure landed last — **silently wrong tensors, not a crash**. Found empirically: new fan-out unit tests reported 6 tasks run when 5 were dispatched, and separately 0, depending on what else in the 1350-test binary was dispatching; `--test-threads=1` passes. A per-module CAS guard is *not* sufficient because the 11 existing dispatch sites do not take it. Latent today (the decode engine dispatches inline on one thread, and the plugin EP calls `disable_persistent_decode_pool()`), live as soon as two sessions in one process both run decode GEMM — which is exactly the concurrency the paged-KV work has to preserve. **(b) Two spinning pools oversubscribe.** With the prototype on and an ORT session alternating in-process, t=16 is bimodal — 0.030 ms on one run, 0.441 ms on the next, same binary, same command. `ONNX_GENAI_CPU_DECODE_THREADS=16` confines the process to CPUs 0–15, which on this host is **8 physical cores** (siblings are adjacent). A blocktime sweep over 500/200/50/20/5/0 µs gave ratios from 0.233 to 72.3 with no monotone trend. The section also records that the compact mask is *correct* and must not be "fixed" to spread — one thread per core across 0,2,…,30 straddles two CCXs and measured worse (0.133 vs 0.079 ms). ### §26.3 — the consequence for the work queue Softmax, Transpose, Concat and the MoE cells should not be attacked as separate kernel problems until (a) is fixed; their t=8/16 numbers are mostly measuring thread wake-up, so per-kernel tuning against that column will have to be redone. The stated exception is §11.5's paged/appendable KV cache, which is a real algorithmic cost no scheduler change touches. ## Scope Documentation only — one file, append-only. No code, no test, no behaviour change. Co-authored-by: Leon <leon@squad.local> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
🔴 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 #1175 +/- ##
=========================================
+ Coverage 0 80.38% +80.38%
=========================================
Files 0 357 +357
Lines 0 154551 +154551
Branches 0 154551 +154551
=========================================
+ Hits 0 124235 +124235
- Misses 0 25764 +25764
- Partials 0 4552 +4552
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
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
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.
What
Two changes to the CPU EP's
RotaryEmbeddingkernel:rotary_units_per_task— Rayon tasks are now sized by work, not by layout. The kernel used to fan out one task per layout unit: one[B,H,S,D]plane or one[B,S,H·D]row. Onrope_llama3_s128that is 4096 elements (16 KiB) per task, and the scheduling cost swamped the work. Tasks now aggregate whole units until each clears 16 Ki elements, targeting ~4 tasks per worker for balance, and collapse to a single serial pass when the pool has one worker.rotate_split_half,rotate_interleaved). The inner loop was left to the autovectoriser, which did not take it.Plus a test-only change: the A/B transform grid now generates
interleaved=1(GPT-J) RoPE models, which were previously uncovered.Why
The old fan-out made the parallel path slower than the serial path at every thread count. Measured on
rope_llama3_s128,exec_kernel.computep50, native-only:Fanning 128 tasks out over a one-worker pool alone cost 36 µs of pure Rayon bookkeeping.
After both changes, same measurement: 242 → 85 µs (t=1), 174 → 67 µs (t=8), 138 → 59 µs (t=16).
Why the vector arms avoid FMA
_mm256_fmsub_pswould drop the intermediate rounding ofcos·x1, so the 8-wide body and the scalar tail would disagree in the last ulp on the same tensor. The same graph would then produce different bits depending on how a thread count happened to split it, which flips the A/B parity check on scheduling alone. Separate_mm256_mul_ps+_mm256_sub_psis bit-identical to the scalar loop. This kernel moves 8 bytes per 3 flops — it is bandwidth-bound, so fusing the multiply is not measurable anyway.Measured: paired interleaved A/B vs ORT 1.27 CPU EP
One driver invocation,
scripts/ort_ab/ab.py, 7 trials × 80 runs × 20 warmups per cell,--threadssets both--native-threadsand--ort-intra-threads.base= merge-base8219fa250,new= this branch. Ratio is native/ORT p50, lower is better; bracket is the min–max band across trials.21 / 21 cells improved. 0 regressed. 5 cells now beat ORT outright. Parity 294/294 PASS.
This does not finish RoPE
Five cells win; sixteen still lose. The remaining loss is not in this kernel — it is thread-pool wake latency, and it is the same root cause behind every "wins at t=1, loses at t=8/16" row in the master benchmark doc (§23.7: Softmax, Transpose, Concat, MoE all have that shape).
Isolated measurement, 16 workers, one 524 288-element elementwise region that costs 388 µs on one thread:
A 20 µs gap — less than the serial glue between two graph nodes — turns a 5.8× speedup into a 1.7× slowdown. Rayon parks its workers at the end of every region; ORT's intra-op pool spins first (
ALLOW_INTRA_OP_SPINNING, default on), so its workers are still hot when the next op arrives and ours are not.exec_kernel.computeforrope_llama3_s128at t=16 is 59 µs against ORT's ~30 µs, but the paired wall is 189 µs — the difference is wake latency, not arithmetic.This repo already has the right mechanism:
decode_spmd's persistent pool with a boundedKMP_BLOCKTIME-style active-spin window. It is wired only to the decode GEMM path. A prototype routing this kernel's fan-out through it takesrope_llama3_s128t=16 from 79 µs to 25 µs native-only, past ORT's 30 µs — but it oversubscribes when an ORT session shares the process, so it needs SMT-aware worker placement to land. That is the follow-up, tracked separately; this PR is the kernel half and stands on its own.Tests
New, each falsified before being kept:
avx2_split_half_is_bit_identical_to_the_scalar_loop— sweepshalfacross 1…128 so the tail runs at every residue mod 8; comparesto_bits(), not a tolerance. Falsified: swapping one_mm256_mul_ps/_mm256_sub_pspair for_mm256_fmsub_psfails it athalf=16.avx2_interleaved_is_bit_identical_to_the_scalar_loop— the de-interleave uses a lane permutation that is its own inverse; a wrong index vector still yields plausible magnitudes, so this compares every element. Falsified: replacing the permutation with the identity fails it athalf=8.task_chunking_covers_every_unit_exactly_once— the derived tiling must be a partition for 4 × 4 × 5 shape/pool combinations, because it feedspar_chunks_mutwhere an off-by-one silently drops or double-writes a row.a_single_worker_pool_never_splits_the_rotationtasks_clear_the_minimum_work_bar_when_the_pool_is_wide— falsified: removing the minimum-task-size floor fails it with "task of 8192 elements is below the 16384-element bar".Existing suites, all green after merging
origin/main(0bd8397f7):cargo test -p onnx-runtime-ep-cpu --lib→ 1328 passed, 0 failedcargo test -p onnx-runtime-ep-cpu-plugin --test plugin_ort_e2e→ 53 passed, 0 failed (real ORT 1.27 session,session.disable_cpu_ep_fallback=1)cargo fmt --all --check,cargo clippy -p onnx-runtime-ep-cpu --all-targets→ cleanNo defer
No node is handed to ORT's CPU EP by this change. The claim policy is untouched;
no_supported_node_is_ever_left_to_the_ort_cpu_epandevery_fixture_loads_with_cpu_fallback_disabledstill pass.Host
AMD EPYC 9V74, 16 physical cores / 32 threads, one NUMA node, 32 MiB shared L3, AVX2+FMA (no AVX-512). Shared and contended, so only paired ratios from a single driver invocation are quoted; absolutes drift by >4×.