From a840ccb0b308fef90b684c5981524fbaa2c64c02 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 04:37:26 +0000 Subject: [PATCH 1/5] perf(cpu): route Transpose and the elementwise fallback through the runtime Two more callers off Rayon and onto the CPU task runtime. Transpose split its output into `units / workers` bands -- one static share per worker, so the whole copy waited for whichever band was unlucky. It now hands the runtime chunks of at least MIN_PARALLEL_TASK_BYTES (64 KiB, a quarter of the whole-operation bar, still comfortably bandwidth-bound) and lets the runtime claim them dynamically and cap the task count. All three split sites -- `block_move`, `blocked_2d`, `batched_blocked_2d` -- use `chunk_runs_mut`, so a task still gets one contiguous band in one call and none of the three bodies changes shape. `simd_activations` already borrowed ORT's pool inside the plugin EP. Its *other* arm -- a native session, or an ORT session whose intra-op pool is one thread, where `prefer_host` correctly declines -- still went to Rayon and still paid the park latency this whole exercise is about: 67us for an isolated region back-to-back, 226us when the previous region ended 20us earlier, which is exactly the spacing between two activations in a decode step. That arm now goes to the native pool. The new tail is outlined in `run_on_task_runtime` rather than written inline, for the reason recorded on `try_host`: growing `run_chunked` repartitions codegen units, and did cost `Relu` 34% at 1 Mi once already with no path change at all. The nesting guards now test `task_runtime::in_task()` as well as `rayon::current_thread_index()`, because both pools can be the outer region while the crate is mid-migration. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/kernels/simd_activations.rs | 53 ++++++++++----- .../src/kernels/transpose.rs | 67 ++++++++++++------- 2 files changed, 78 insertions(+), 42 deletions(-) diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs index b612487108..e77687b968 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/simd_activations.rs @@ -570,24 +570,49 @@ where body(input, output); return; } - // Nesting a `par_chunks` inside an outer parallel region only adds - // scheduling overhead: the outer region is already keeping the pool busy. - if rayon::current_thread_index().is_some() { + // Nesting a split inside an outer parallel region only adds scheduling + // overhead: the outer region is already keeping the pool busy. Rayon is + // still checked because plenty of this crate fans out through it directly. + if rayon::current_thread_index().is_some() || crate::task_runtime::in_task() { body(input, output); return; } - let Some(chunk) = par_chunk_len(len, rayon::current_num_threads()) else { + let Some(chunk) = par_chunk_len(len, crate::task_runtime::width()) else { body(input, output); return; }; note_parallel_dispatch(); + run_on_task_runtime(input, output, chunk, &body); +} - use rayon::prelude::*; - output - .par_chunks_mut(chunk) - .zip(input.par_chunks(chunk)) - .for_each(|(o, i)| body(i, o)); +/// Splits `input`/`output` into `chunk`-sized pieces across the native task +/// runtime and runs `body` on each. +/// +/// Bit-identical to calling `body` on the whole slice, for the same reason +/// [`run_on_host`] is: the kernels are elementwise, every chunk is a multiple of +/// eight lanes, and none is shorter than [`PAR_MIN_CHUNK`], so no chunk takes a +/// different code path from its neighbours. +/// +/// The runtime's pool replaces rayon here because rayon parks its workers +/// between regions and an elementwise op is *all* region boundary: an isolated +/// region measured 67 us back-to-back and 226 us when the previous one ended +/// 20 us earlier, which is the spacing between two activations in a decode step. +/// +/// `#[inline(never)]` for the same reason as [`try_host`]: `run_chunked`'s +/// callers are the hottest elementwise kernels in the crate, and growing it has +/// repartitioned codegen units before -- 34% on `Relu` at 1 Mi with no path +/// change at all. +#[inline(never)] +fn run_on_task_runtime(input: &[f32], output: &mut [f32], chunk: usize, body: &F) +where + F: Fn(&[f32], &mut [f32]) + Sync, +{ + debug_assert_eq!(input.len(), output.len()); + crate::task_runtime::chunk_runs_mut(output, chunk, 1, |first, out| { + let start = first * chunk; + body(&input[start..start + out.len()], out); + }); } #[cfg(test)] @@ -675,20 +700,16 @@ where body(input, output); return; } - if rayon::current_thread_index().is_some() { + if rayon::current_thread_index().is_some() || crate::task_runtime::in_task() { body(input, output); return; } - let Some(chunk) = par_chunk_len_rows(len, width, rayon::current_num_threads()) else { + let Some(chunk) = par_chunk_len_rows(len, width, crate::task_runtime::width()) else { body(input, output); return; }; - use rayon::prelude::*; - output - .par_chunks_mut(chunk) - .zip(input.par_chunks(chunk)) - .for_each(|(o, i)| body(i, o)); + run_on_task_runtime(input, output, chunk, &body); } // --------------------------------------------------------------------------- diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs b/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs index 98cd1c20df..bc42aed295 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs @@ -76,6 +76,17 @@ const TILE: usize = 64; /// more cores start helping. const MIN_PARALLEL_BYTES: usize = 256 * 1024; +/// Minimum output bytes in one task's chunk. +/// +/// A quarter of [`MIN_PARALLEL_BYTES`]: the whole-operation bar exists because +/// below it the copy finishes faster than a fan-out can be scheduled, and a +/// single chunk faces the same arithmetic against the task runtime's much +/// cheaper per-task cost (a `fetch_sub` and an indirect call, not a fork/join). +/// 64 KiB is still comfortably past the point where a copy is bandwidth-bound, +/// so a chunk this size scales, while being small enough that a 1 MiB transpose +/// has 16 chunks to balance across rather than one per thread. +const MIN_PARALLEL_TASK_BYTES: usize = 64 * 1024; + /// Dtype-generic Transpose kernel carrying the resolved `perm`. pub struct TransposeKernel { /// Axis permutation; `None` means reverse all axes. @@ -389,11 +400,10 @@ fn block_move(src: &[u8], dst: &mut [u8], shape: &[usize], perm: &[usize], block }; match parallel_blocks_per_task(blocks, block) { - Some(per_task) => { - use rayon::prelude::*; - dst.par_chunks_mut(per_task * block) - .enumerate() - .for_each(|(task, chunk)| body(chunk, task * per_task)); + Some(per_chunk) => { + crate::task_runtime::chunk_runs_mut(dst, per_chunk * block, 1, |first, run| { + body(run, first * per_chunk); + }); } None => body(dst, 0), } @@ -407,13 +417,10 @@ fn blocked_2d(src: &[u8], dst: &mut [u8], rows: usize, cols: usize, esize: usize // worker a `&mut [u8]` without any unsafe pointer splitting. let dst_row = rows * esize; match parallel_blocks_per_task(cols, dst_row) { - Some(per_task) => { - use rayon::prelude::*; - dst.par_chunks_mut(per_task * dst_row) - .enumerate() - .for_each(|(task, chunk)| { - blocked_2d_band(src, chunk, rows, cols, esize, task * per_task); - }); + Some(per_chunk) => { + crate::task_runtime::chunk_runs_mut(dst, per_chunk * dst_row, 1, |first, run| { + blocked_2d_band(src, run, rows, cols, esize, first * per_chunk); + }); } None => blocked_2d_band(src, dst, rows, cols, esize, 0), } @@ -469,16 +476,13 @@ fn batched_blocked_2d( plane: usize, ) { match parallel_blocks_per_task(batch, plane) { - Some(per_task) => { - use rayon::prelude::*; - dst.par_chunks_mut(per_task * plane) - .enumerate() - .for_each(|(task, chunk)| { - for (n, dst_plane) in chunk.chunks_mut(plane).enumerate() { - let base = (task * per_task + n) * plane; - blocked_2d_band(&src[base..base + plane], dst_plane, rows, cols, esize, 0); - } - }); + Some(per_chunk) => { + crate::task_runtime::chunk_runs_mut(dst, per_chunk * plane, 1, |first, run| { + for (n, dst_plane) in run.chunks_mut(plane).enumerate() { + let base = (first * per_chunk + n) * plane; + blocked_2d_band(&src[base..base + plane], dst_plane, rows, cols, esize, 0); + } + }); } None => { // Too few planes to split, so let each plane decide for itself @@ -524,16 +528,27 @@ fn unflatten(mut flat: usize, shape: &[usize]) -> Vec { idx } -/// `Some(units_per_task)` when splitting `units` of `unit_bytes` each pays off. +/// `Some(units_per_chunk)` when splitting `units` of `unit_bytes` each pays off. +/// +/// The chunk is a *balance* unit, not a per-thread share. Sizing it as +/// `units / workers` gives every worker one static band and makes the whole +/// transpose wait for whichever band was unlucky -- a descheduled thread, an SMT +/// sibling, a page fault on first touch. The task runtime claims chunks +/// dynamically and caps the task count itself, so handing it several chunks per +/// thread costs nothing on an idle machine and absorbs a straggler on a busy +/// one. +/// +/// [`MIN_PARALLEL_TASK_BYTES`] keeps a chunk from falling below the size where +/// a copy is bandwidth-bound rather than latency-bound; below that the +/// per-chunk bookkeeping starts to show. fn parallel_blocks_per_task(units: usize, unit_bytes: usize) -> Option { if units.saturating_mul(unit_bytes) < MIN_PARALLEL_BYTES { return None; } - let workers = rayon::current_num_threads().max(1); - if workers < 2 || units < 2 { + if crate::task_runtime::width() < 2 || units < 2 { return None; } - Some(units.div_ceil(workers).max(1)) + Some(MIN_PARALLEL_TASK_BYTES.div_ceil(unit_bytes.max(1)).clamp(1, units)) } #[cfg(test)] From a556a94109d92bf8d2f3d6a28cbf47cbca92489d Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 04:57:13 +0000 Subject: [PATCH 2/5] bench: generate isolated activation cells for the A/B harness The elementwise ops are the ones most exposed to scheduling overhead -- a Gelu over an MLP intermediate is a few hundred microseconds of arithmetic wrapped around a fan-out -- so they are the sharpest test of a task runtime, and the harness had no cells for them. The sweep is by token count at fixed hidden/intermediate widths, which walks the tensor from L2-resident through L3 to memory-resident while holding the fan-out cost fixed. That is the axis that decides whether parallelising is worth it at all, so it is the axis a scheduler change has to be measured along. FastGelu rather than Gelu because Gelu only joins the default domain at opset 20, and FastGelu is what a real transformer graph carries here. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- scripts/ort_ab/gen_activations.py | 109 ++++++++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 scripts/ort_ab/gen_activations.py diff --git a/scripts/ort_ab/gen_activations.py b/scripts/ort_ab/gen_activations.py new file mode 100644 index 0000000000..59573847fb --- /dev/null +++ b/scripts/ort_ab/gen_activations.py @@ -0,0 +1,109 @@ +#!/usr/bin/env python3 +"""Isolated single-node graphs for the elementwise activations. + +These are the cheapest ops in a transformer and the ones most exposed to +scheduling overhead: a Gelu over an MLP intermediate tensor is a few hundred +microseconds of arithmetic wrapped around a fan-out, so the fan-out is a +first-order cost rather than a rounding error. That makes them the sharpest +test of a task runtime. + +Shapes are MLP-intermediate and hidden-state shaped, taken from public model +configs, and swept over a size range that brackets the point where parallelism +starts to pay: L2-resident, L3-resident, and memory-resident. + +SYNTHETIC DATA NOTICE: no trained weights. Only the hidden / intermediate / +sequence dimensions come from public model configs; tensor contents are the +benchmark harness's deterministic synthetic pattern, fed identically to both +runtimes. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import onnx +from onnx import TensorProto, helper + +OPSET = 17 + + +def _save(path: Path, graph: onnx.GraphProto, *, ms_domain: bool = False) -> None: + opsets = [helper.make_opsetid("", OPSET)] + if ms_domain: + opsets.append(helper.make_opsetid("com.microsoft", 1)) + model = helper.make_model(graph, opset_imports=opsets) + model.ir_version = 10 + onnx.checker.check_model(model) + path.parent.mkdir(parents=True, exist_ok=True) + onnx.save(model, str(path)) + print(f"wrote {path}") + + +def build_unary(path: Path, op: str, *, rows: int, cols: int, ms_domain: bool = False) -> None: + """A single unary elementwise node over a [rows, cols] float tensor.""" + domain = "com.microsoft" if ms_domain else "" + node = helper.make_node(op, ["x"], ["y"], domain=domain) + graph = helper.make_graph( + [node], + op.lower(), + [helper.make_tensor_value_info("x", TensorProto.FLOAT, [rows, cols])], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [rows, cols])], + ) + _save(path, graph, ms_domain=ms_domain) + + +def build_binary(path: Path, op: str, *, rows: int, cols: int) -> None: + """A single binary elementwise node over two [rows, cols] float tensors.""" + node = helper.make_node(op, ["x", "w"], ["y"]) + graph = helper.make_graph( + [node], + op.lower(), + [ + helper.make_tensor_value_info("x", TensorProto.FLOAT, [rows, cols]), + helper.make_tensor_value_info("w", TensorProto.FLOAT, [rows, cols]), + ], + [helper.make_tensor_value_info("y", TensorProto.FLOAT, [rows, cols])], + ) + _save(path, graph) + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--out", type=Path, required=True) + args = ap.parse_args() + out = args.out + + # --- The size sweep ------------------------------------------------- + # llama3-8B hidden is 4096, intermediate is 14336. One row per token, so + # the token count is what moves the tensor through the cache hierarchy: + # 32 tokens is 512 KiB (L2), 512 tokens is 8 MiB (L3), 4096 tokens is + # 64 MiB (memory). The fan-out cost is fixed, so the ratio between it and + # the body is what changes across this sweep. + # Gelu is only in the default domain from opset 20; the contrib FastGelu + # is what a real transformer graph carries at this opset, and it is the + # one the CPU EP fuses. + for tokens in (1, 32, 512, 4096): + build_unary( + out / f"act_fastgelu_mlp_t{tokens}.onnx", + "FastGelu", + rows=tokens, + cols=14336, + ms_domain=True, + ) + for tokens in (32, 512, 4096): + build_unary(out / f"act_relu_hidden_t{tokens}.onnx", "Relu", rows=tokens, cols=4096) + build_unary(out / f"act_sigmoid_hidden_t{tokens}.onnx", "Sigmoid", rows=tokens, cols=4096) + build_unary(out / f"act_tanh_hidden_t{tokens}.onnx", "Tanh", rows=tokens, cols=4096) + + # --- The binary elementwise ops that surround them ------------------ + # SwiGLU's gate multiply and the residual add are the same fan-out with a + # second input stream, which doubles the bandwidth per unit of arithmetic. + for tokens in (32, 512, 4096): + build_binary(out / f"act_mul_gate_t{tokens}.onnx", "Mul", rows=tokens, cols=14336) + for tokens in (512, 4096): + build_binary(out / f"act_add_residual_t{tokens}.onnx", "Add", rows=tokens, cols=4096) + + +if __name__ == "__main__": + main() From b75da841dcc5cf81c09e92e6e74470d1f08ccac0 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 07:07:53 +0000 Subject: [PATCH 3/5] docs(bench): record the transpose and activation matrices, the controls, and the GEMM finding Adds scripts/ort_ab/gen_gemm.py, which generates the MatMulNBits and dense MatMul cells the control grid needed. That control was meant to be a null result and instead found the largest remaining loss in the CPU EP: decode-shaped quantised GEMM is 4.6x slower at 32 threads than at 16. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../2026-08-15-cpu-ep-vs-ort-attention-moe.md | 57 ++++++++ scripts/ort_ab/gen_gemm.py | 132 ++++++++++++++++++ 2 files changed, 189 insertions(+) create mode 100644 scripts/ort_ab/gen_gemm.py diff --git a/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md b/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md index f05925cb4f..c33e0973fb 100644 --- a/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md +++ b/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md @@ -3061,3 +3061,60 @@ because what it removes is a per-region park, not a per-element cost. The decode-shaped `sm_decode_h32_kv*` cells stay serial and stay flat (18–182 µs, within ±1 µs of base at every width), and `sm_whisper_cross` is neutral at every width — both are the intended outcome, not a miss. +### 31.8 Transpose and the activations, and what the controls say + +Activations, our own native ms, `base>new`, 5 trials × 7 runs: + +| cell | t=1 | t=2 | t=4 | t=8 | t=16 | t=32 | +|---|---|---|---|---|---|---| +| act_relu_hidden_t512 | 0.838>0.822 | 0.546>0.488 | 0.295>0.251 | 0.276>0.183 | 0.381>0.130 | 1.201>0.303 | +| act_tanh_hidden_t512 | 0.804>0.803 | 0.806>0.794 | 0.457>0.417 | 0.419>0.317 | 0.485>0.204 | 1.441>0.386 | +| act_sigmoid_hidden_t512 | 0.826>0.837 | 0.832>0.808 | 0.640>0.559 | 0.462>0.296 | 0.513>0.214 | 1.176>0.437 | +| act_tanh_hidden_t4096 | 6.398>6.542 | 6.331>6.436 | 4.300>3.412 | 3.609>3.494 | 3.691>3.490 | 5.991>2.933 | +| act_sigmoid_hidden_t4096 | 6.532>6.680 | 6.488>6.605 | 3.547>3.465 | 3.560>3.476 | 3.672>3.539 | 6.070>2.356 | +| act_fastgelu_mlp_t512 | 5.466>5.448 | 5.169>5.294 | 2.815>2.765 | 2.156>1.864 | 1.949>1.553 | 4.461>1.463 | + +27 improved cells against 5 regressed at a 5% deadband, and the five are 1–6 µs +absolute or bandwidth-bound cells whose distributions overlap outright. Up to +3.6× at t=16 and 4× at t=32. + +Transpose needed 15 trials rather than 7 before it said anything stable — at 7 +it produced a `tr_bert_b8_s128` "regression" that was entirely dispersion: + +| cell | t=1 | t=2 | t=4 | t=8 | t=16 | t=32 | +|---|---|---|---|---|---|---| +| tr_bert_b8_s128 | 0.117>0.119 | 0.152>0.153 | 0.111>0.112 | 0.134>0.140 | 0.261>0.248 | 2.658>2.670 | +| tr_llama3_s512 | 0.718>0.705 | 0.390>0.389 | 0.227>0.228 | 0.253>0.239 | 0.430>0.352 | 3.412>3.101 | +| tr_whisper_s1500 | 0.566>0.593 | 0.448>0.442 | 0.353>0.358 | 0.307>0.296 | 0.397>0.380 | 3.292>3.017 | + +4 improved, 0 regressed. Transpose is a straight copy, so it is bandwidth-bound +long before it is scheduler-bound; the gain is small and lives at the wide end. + +**The controls.** Two families this work does not touch were run at the same six +widths purely to establish the noise floor: 24 GQA attention cells (32 improved +/ 27 regressed) and 20 GEMM cells (19 / 12). Both are symmetric, which is what +licenses reading the 3–15× numbers as real rather than as luck. Anyone repeating +this on this host should budget for ±25% on a single 7-trial cell and treat any +claim smaller than that as unmeasured. + +### 31.9 The biggest remaining loss is GEMM, and it is the same disease + +The GEMM control was meant to be a null result and instead found the largest +single loss in the CPU EP. Decode-shaped `MatMulNBits`, native ms: + +| cell | t=8 | t=16 | t=32 | ratio vs ORT at t=32 | +|---|---|---|---|---| +| gemm_nbits_qwen3_0p6b_qkv_t8 | 0.67 | 0.64 | **2.98** | 25.99 | +| gemm_nbits_qwen3_0p6b_mlp_t8 | 1.17 | 1.03 | **3.81** | 18.73 | +| gemm_nbits_llama3_8b_qkv_t8 | 4.15 | 3.05 | **6.07** | 8.64 | + +Going from 16 to 32 threads makes these **4.6× slower**, on the operator that +carries most of a decode step's arithmetic. That is the §26 signature exactly, +and `matmul_nbits.rs` is still the largest raw-Rayon fan-out in the crate. + +Prefill-shaped GEMM (`t128`, `t512`) does not show it: those regions are long +enough to amortise a park. It is specifically the small, frequent, decode-shaped +region that the global pool cannot serve, which is the whole thesis of §26. + +Nothing in this phase touches it. It is the next target, and the runtime it +needs already exists. diff --git a/scripts/ort_ab/gen_gemm.py b/scripts/ort_ab/gen_gemm.py new file mode 100644 index 0000000000..b6be574b1d --- /dev/null +++ b/scripts/ort_ab/gen_gemm.py @@ -0,0 +1,132 @@ +#!/usr/bin/env python3 +"""Generate single-node GEMM benchmark cells: block-quantised +`com.microsoft::MatMulNBits` and dense `MatMul`. + +These are the two operators that carry a decode step's arithmetic, and they are +also the largest remaining raw-Rayon fan-outs in the CPU EP. They are generated +here so the task-runtime A/B can show what it does *not* touch as well as what +it does. + +SYNTHETIC DATA NOTICE: no trained weights. Only the hidden/intermediate sizes +come from public model configs; the packed `B` matrix and the scales are +deterministic byte patterns, and the activations come from the benchmark +harness's own synthetic pattern, fed identically to both runtimes. + +`MatMulNBits` slot order (com.microsoft): + 0 A 1 B 2 scales 3 zero_points 4 g_idx 5 bias +`B` is `[n, n_blocks_per_col, blob_size]` uint8 with `blob_size = block_size/2` +for 4-bit, and `scales` is `[n * n_blocks_per_col]`. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import onnx +from onnx import TensorProto, helper, numpy_helper + +OPSET = 17 + +# (name, k, n) -- one projection each from a few public configs. +NBITS_SHAPES = { + "qwen3_0p6b_qkv": (1024, 3072), + "qwen3_0p6b_mlp": (1024, 6144), + "llama3_8b_qkv": (4096, 6144), + "llama3_8b_mlp": (4096, 14336), +} + +# Token counts: 1 = decode, the rest walk prefill up through the caches. +NBITS_TOKENS = (1, 8, 128, 512) + +# Dense f32 MatMul cells, [m, k] x [k, n]. +DENSE_SHAPES = { + "sq_512": (512, 512, 512), + "sq_1024": (1024, 1024, 1024), + "tall_128x4096": (128, 4096, 4096), + "decode_1x4096": (1, 4096, 4096), +} + +BLOCK_SIZE = 32 + + +def build_matmul_nbits(path: Path, *, tokens: int, k: int, n: int) -> None: + blocks = (k + BLOCK_SIZE - 1) // BLOCK_SIZE + blob = BLOCK_SIZE // 2 + + rng = np.random.default_rng(0x5EBA5) + b = rng.integers(0, 256, size=(n, blocks, blob), dtype=np.uint8) + # Small positive scales keep the dequantised weights in a sane range so the + # parity check compares meaningful numbers rather than saturated ones. + scales = (rng.random((n * blocks,), dtype=np.float32) * 0.01 + 0.001).astype(np.float32) + + node = helper.make_node( + "MatMulNBits", + inputs=["A", "B", "scales"], + outputs=["Y"], + domain="com.microsoft", + name="matmul_nbits", + K=k, + N=n, + bits=4, + block_size=BLOCK_SIZE, + accuracy_level=0, + ) + graph = helper.make_graph( + [node], + "matmul_nbits", + [helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, tokens, k])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, tokens, n])], + initializer=[ + numpy_helper.from_array(b, "B"), + numpy_helper.from_array(scales, "scales"), + ], + ) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", OPSET), + helper.make_opsetid("com.microsoft", 1), + ], + ) + model.ir_version = 10 + onnx.save(model, str(path)) + + +def build_matmul(path: Path, *, m: int, k: int, n: int) -> None: + rng = np.random.default_rng(0x5EBA6) + b = (rng.random((k, n), dtype=np.float32) - 0.5).astype(np.float32) + node = helper.make_node("MatMul", inputs=["A", "B"], outputs=["Y"], name="matmul") + graph = helper.make_graph( + [node], + "matmul", + [helper.make_tensor_value_info("A", TensorProto.FLOAT, [m, k])], + [helper.make_tensor_value_info("Y", TensorProto.FLOAT, [m, n])], + initializer=[numpy_helper.from_array(b, "B")], + ) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", OPSET)]) + model.ir_version = 10 + onnx.save(model, str(path)) + + +def main(out: Path) -> None: + out.mkdir(parents=True, exist_ok=True) + made = [] + for name, (k, n) in NBITS_SHAPES.items(): + for tokens in NBITS_TOKENS: + path = out / f"gemm_nbits_{name}_t{tokens}.onnx" + build_matmul_nbits(path, tokens=tokens, k=k, n=n) + made.append(path) + for name, (m, k, n) in DENSE_SHAPES.items(): + path = out / f"gemm_dense_{name}.onnx" + build_matmul(path, m=m, k=k, n=n) + made.append(path) + for p in made: + print(p) + + +if __name__ == "__main__": + ap = argparse.ArgumentParser() + ap.add_argument("--out", type=Path, default=Path(__file__).resolve().parent / "models" / "gemm") + main(ap.parse_args().out) From 71a33681e71c4d4e4e83dba608a3edcd1b88f7b7 Mon Sep 17 00:00:00 2001 From: Justin Chu Date: Tue, 18 Aug 2026 07:35:13 +0000 Subject: [PATCH 4/5] style(cpu): rustfmt the transpose task-runtime call sites Pure rustfmt output, no behaviour change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs | 6 +++++- docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md | 1 + 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs b/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs index bc42aed295..271303df3b 100644 --- a/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs +++ b/crates/onnx-runtime-ep-cpu/src/kernels/transpose.rs @@ -548,7 +548,11 @@ fn parallel_blocks_per_task(units: usize, unit_bytes: usize) -> Option { if crate::task_runtime::width() < 2 || units < 2 { return None; } - Some(MIN_PARALLEL_TASK_BYTES.div_ceil(unit_bytes.max(1)).clamp(1, units)) + Some( + MIN_PARALLEL_TASK_BYTES + .div_ceil(unit_bytes.max(1)) + .clamp(1, units), + ) } #[cfg(test)] diff --git a/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md b/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md index c33e0973fb..4b763ca686 100644 --- a/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md +++ b/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md @@ -3061,6 +3061,7 @@ because what it removes is a per-region park, not a per-element cost. The decode-shaped `sm_decode_h32_kv*` cells stay serial and stay flat (18–182 µs, within ±1 µs of base at every width), and `sm_whisper_cross` is neutral at every width — both are the intended outcome, not a miss. + ### 31.8 Transpose and the activations, and what the controls say Activations, our own native ms, `base>new`, 5 trials × 7 runs: From 064afde33cd1fa2e74421143b412a0c9d225dcb8 Mon Sep 17 00:00:00 2001 From: justinchuby <223556219+Copilot@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:28:54 -0700 Subject: [PATCH 5/5] fix(cpu): construct the slot-exhaustion precondition and de-collide doc numbering The slot-exhaustion pool test slept a fixed 50 ms and assumed the OS had scheduled all SLOT_COUNT holder threads into their dispatch bodies inside that window. Under a contended box (the rest of the lib suite running in parallel) some holders had not yet claimed a slot, so the probe found a free slot, succeeded, and the "declines instead of blocking" assertion flaked. Replace the sleep with an explicit wait on a per-holder `held` counter so the probe only runs once every slot is provably occupied, bounded by a 10 s deadline so a starved runner fails loudly instead of hanging. Verified green under 20 background CPU spinners. Also renumber the task-runtime doc subsections that were still labelled 31.7/31.8/31.9 -- stale from before the section became 33 -- which collided with the genuine section 31.7/31.8. They are now 33.7/33.8/33.9. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../src/task_runtime/pool.rs | 32 +++++++++++++++++-- .../2026-08-15-cpu-ep-vs-ort-attention-moe.md | 10 +++--- 2 files changed, 34 insertions(+), 8 deletions(-) diff --git a/crates/onnx-runtime-ep-cpu/src/task_runtime/pool.rs b/crates/onnx-runtime-ep-cpu/src/task_runtime/pool.rs index a84b39ad8a..8dfda11725 100644 --- a/crates/onnx-runtime-ep-cpu/src/task_runtime/pool.rs +++ b/crates/onnx-runtime-ep-cpu/src/task_runtime/pool.rs @@ -736,21 +736,47 @@ mod tests { let pool = Arc::new(TaskPool::new(4)); let hold = Arc::new(AtomicBool::new(true)); let declined = Arc::new(AtomicUsize::new(0)); + // Count how many holder dispatches have actually entered their body, so + // the probe below runs only once every slot is provably occupied. + // Sleeping a fixed 50 ms instead assumes the OS scheduled all + // `SLOT_COUNT` holder threads inside that window -- which fails on a + // contended box (e.g. the rest of the test suite running in parallel), + // leaving slots free, the probe succeeding, and the assertion flaking. + // Constructing the precondition removes the race rather than papering + // it over. + let held = Arc::new(AtomicUsize::new(0)); thread::scope(|scope| { for _ in 0..SLOT_COUNT { let pool = Arc::clone(&pool); let hold = Arc::clone(&hold); + let held = Arc::clone(&held); scope.spawn(move || { + // One increment per holder even though the fan-out has two + // tasks and a pool worker may run the second. + let counted = AtomicBool::new(false); let _ = pool.dispatch(2, &|_| { + if !counted.swap(true, Ordering::Relaxed) { + held.fetch_add(1, Ordering::Release); + } while hold.load(Ordering::Relaxed) { std::hint::spin_loop(); } }); }); } - // Give the holders time to take every slot, then prove the next - // dispatcher is declined rather than blocked. - thread::sleep(Duration::from_millis(50)); + // Wait until every holder is provably inside its body -- i.e. every + // slot is claimed -- then prove the next dispatcher is declined + // rather than blocked. A generous deadline keeps a starved runner + // from hanging the suite forever. + let deadline = Instant::now() + Duration::from_secs(10); + while held.load(Ordering::Acquire) < SLOT_COUNT && Instant::now() < deadline { + std::hint::spin_loop(); + } + assert_eq!( + held.load(Ordering::Acquire), + SLOT_COUNT, + "not every holder took a slot within the deadline" + ); for _ in 0..4 { if !pool.dispatch(2, &|_| {}) { declined.fetch_add(1, Ordering::Relaxed); diff --git a/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md b/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md index 4b763ca686..ad787b75c3 100644 --- a/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md +++ b/docs/benchmarks/2026-08-15-cpu-ep-vs-ort-attention-moe.md @@ -2973,7 +2973,7 @@ path reads an environment variable for test purposes, and the tests do not race on a global by construction: `isolated_pool` builds a private pool rather than reconfiguring the shared one. -### 31.7 RoPE and Softmax, moved onto it +### 33.7 RoPE and Softmax, moved onto it Real ORT A/B, one driver invocation per pair, 7 trials × 7 runs. Ratio is native ÷ ORT, lower is better, `base>new`: @@ -3009,7 +3009,7 @@ The decode-shaped cells are deliberately unmoved: `rope_*_s1` and arms, flat to within noise at every width. A one-token RoPE should not be fanning out and still does not. -### 31.7.1 The scar: an unrelated edit de-vectorised the activations +### 33.7.1 The scar: an unrelated edit de-vectorised the activations Rewiring RoPE and Softmax made **`Tanh` and `Sigmoid` 2× slower at one thread**, in a file the change did not touch, on inputs that never reach a parallel path. @@ -3037,7 +3037,7 @@ Two things about this are worth keeping: neither helper carries the latter; every caller is already inside an `avx2,fma` function, so the features arrive with the inline. -### 31.7.2 Re-verified against §29's Softmax, not the one it replaced +### 33.7.2 Re-verified against §29's Softmax, not the one it replaced The Softmax numbers above were taken before #1219 landed, so their baseline was a `softmax_rows` that still copied `src` into `dst` before reducing. Merging @@ -3062,7 +3062,7 @@ decode-shaped `sm_decode_h32_kv*` cells stay serial and stay flat (18–182 µs, within ±1 µs of base at every width), and `sm_whisper_cross` is neutral at every width — both are the intended outcome, not a miss. -### 31.8 Transpose and the activations, and what the controls say +### 33.8 Transpose and the activations, and what the controls say Activations, our own native ms, `base>new`, 5 trials × 7 runs: @@ -3098,7 +3098,7 @@ licenses reading the 3–15× numbers as real rather than as luck. Anyone repeat this on this host should budget for ±25% on a single 7-trial cell and treat any claim smaller than that as unmeasured. -### 31.9 The biggest remaining loss is GEMM, and it is the same disease +### 33.9 The biggest remaining loss is GEMM, and it is the same disease The GEMM control was meant to be a null result and instead found the largest single loss in the CPU EP. Decode-shaped `MatMulNBits`, native ms: