diff --git a/crates/onnx-runtime-ep-cpu/benches/activation_bench.rs b/crates/onnx-runtime-ep-cpu/benches/activation_bench.rs index 206ac68883..41152162ce 100644 --- a/crates/onnx-runtime-ep-cpu/benches/activation_bench.rs +++ b/crates/onnx-runtime-ep-cpu/benches/activation_bench.rs @@ -79,6 +79,9 @@ fn time(name: &str, kernel: &dyn Kernel, ins: &[TensorView], out: &mut Tensor, e } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + // Decode (`[1, 1, H]`), prefill (`[1, S, H]`) and two sizes around the // vector dispatch threshold. let shapes: [(&str, Vec); 5] = [ diff --git a/crates/onnx-runtime-ep-cpu/benches/common/mod.rs b/crates/onnx-runtime-ep-cpu/benches/common/mod.rs index bfd6ae2884..146b0120c1 100644 --- a/crates/onnx-runtime-ep-cpu/benches/common/mod.rs +++ b/crates/onnx-runtime-ep-cpu/benches/common/mod.rs @@ -203,3 +203,30 @@ pub fn assert_close(actual: &[f32], expected: &[f32], tolerance: f32) { ); } } + +/// Put this benchmark process into the same decode thread topology a served +/// session runs in. +/// +/// `CpuExecutionProvider::initialize` is the earliest per-session hook and the +/// only place an explicit `ONNX_GENAI_CPU_DECODE_THREADS` budget becomes a +/// process-wide bound on prefill/MLAS Rayon parallelism and, on Linux, on CPU +/// affinity. A benchmark that never calls it runs every row on an *unbounded* +/// process, so a `t=N` row measures N decode workers competing with a +/// full-width Rayon pool on every core -- not the configuration production runs +/// (#1749). +/// +/// Deliberately the real `CpuExecutionProvider::initialize` rather than a copy +/// of what it currently does. A reimplementation would silently stop matching +/// the moment `initialize` grows a second responsibility, which is exactly the +/// drift that leaves a benchmark measuring a configuration nothing ships. +/// +/// A no-op unless a budget is set, and idempotent: the underlying bound latches +/// on first call. Benches that install their own fixed-width Rayon pool (see +/// `kernels.rs`, which sweeps `[1, 8]`) will oversubscribe that pool onto a +/// smaller budget's cores -- which is what production does with the same two +/// settings, and is why this is applied there too rather than special-cased. +pub fn init_decode_topology() { + CpuExecutionProvider::new() + .initialize(&Default::default()) + .expect("the CPU EP must initialize"); +} diff --git a/crates/onnx-runtime-ep-cpu/benches/gqa_decode.rs b/crates/onnx-runtime-ep-cpu/benches/gqa_decode.rs index b3dcd649f1..37f936e4a3 100644 --- a/crates/onnx-runtime-ep-cpu/benches/gqa_decode.rs +++ b/crates/onnx-runtime-ep-cpu/benches/gqa_decode.rs @@ -71,6 +71,10 @@ fn bind_decode_pool_width(threads: usize) { /// `sdpa_decode_group` call over the same KV window. Single-threaded on purpose /// — this isolates the memory-traffic difference from the decode scheduler. fn bench_sdpa_group_vs_row(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group_bench = c.benchmark_group("sdpa_group_vs_row"); for (label, num_heads, kv_num_heads, head_size) in GEOMETRIES { let group = num_heads / kv_num_heads; @@ -217,6 +221,10 @@ fn gqa_kernel( /// the bench twice (once with the variable set to `0`, once to `1`) to get the /// A/B. The current setting is reported in the benchmark id. fn bench_gqa_kernel_decode(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let fused = std::env::var("ONNX_GENAI_GQA_GROUP_FUSED").unwrap_or_else(|_| "default".into()); let mut group_bench = c.benchmark_group("gqa_kernel_decode"); // 8 workers is the bounded decode-pool width the runtime uses for M=1. @@ -266,6 +274,10 @@ fn bench_gqa_kernel_decode(c: &mut Criterion) { /// read either way, so the ratio between arms reads directly as the fraction of /// cache traffic the append removes. fn bench_gqa_half_kv_append(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group_bench = c.benchmark_group("gqa_half_kv_append"); bind_decode_pool_width(8); for (label, num_heads, kv_num_heads, head_size) in GEOMETRIES { diff --git a/crates/onnx-runtime-ep-cpu/benches/half_decode_gemv_ab.rs b/crates/onnx-runtime-ep-cpu/benches/half_decode_gemv_ab.rs index 98e898a9c2..1d1c6777f6 100644 --- a/crates/onnx-runtime-ep-cpu/benches/half_decode_gemv_ab.rs +++ b/crates/onnx-runtime-ep-cpu/benches/half_decode_gemv_ab.rs @@ -144,6 +144,9 @@ fn max_relative_deviation(got: &[f32], want: &[f32]) -> f32 { } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + // Shapes a decode step actually issues: a 7B-class MLP projection, a // vocabulary head (the widest weight in the graph), and a square control. let shapes: Vec<(&str, usize, usize)> = match std::env::var("PROBE_SHAPE").as_deref() { diff --git a/crates/onnx-runtime-ep-cpu/benches/half_prefill_route_ab.rs b/crates/onnx-runtime-ep-cpu/benches/half_prefill_route_ab.rs index dcb9ee051e..94e969e911 100644 --- a/crates/onnx-runtime-ep-cpu/benches/half_prefill_route_ab.rs +++ b/crates/onnx-runtime-ep-cpu/benches/half_prefill_route_ab.rs @@ -105,6 +105,9 @@ fn build_kernel(dtype: FloatDType, m: usize, k: usize, n: usize) -> Box = match std::env::var("PROBE_SHAPE").as_deref() { Ok("big") => vec![(4096, 11008)], Ok("small") => vec![(2048, 2048)], diff --git a/crates/onnx-runtime-ep-cpu/benches/int4_acc0_attribution.rs b/crates/onnx-runtime-ep-cpu/benches/int4_acc0_attribution.rs index 85ba498e6b..14a4dcee96 100644 --- a/crates/onnx-runtime-ep-cpu/benches/int4_acc0_attribution.rs +++ b/crates/onnx-runtime-ep-cpu/benches/int4_acc0_attribution.rs @@ -52,6 +52,8 @@ //! cargo bench -p onnx-runtime-ep-cpu --bench int4_acc0_attribution //! ``` +mod common; + use std::time::Instant; /// Shapes taken from the projections that dominate llama/qwen decode. @@ -207,6 +209,9 @@ fn bench(label: &str, reps: usize, mut f: impl FnMut()) -> f64 { } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + let reps: usize = std::env::var("REPS") .ok() .and_then(|v| v.parse().ok()) diff --git a/crates/onnx-runtime-ep-cpu/benches/int4_decode_loop_ab.rs b/crates/onnx-runtime-ep-cpu/benches/int4_decode_loop_ab.rs index 1945747df1..35275720a0 100644 --- a/crates/onnx-runtime-ep-cpu/benches/int4_decode_loop_ab.rs +++ b/crates/onnx-runtime-ep-cpu/benches/int4_decode_loop_ab.rs @@ -179,6 +179,9 @@ struct Weight { } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + let block_size: usize = std::env::var("PROBE_BLOCK") .ok() .and_then(|v| v.parse().ok()) diff --git a/crates/onnx-runtime-ep-cpu/benches/int4_prefill_route_ab.rs b/crates/onnx-runtime-ep-cpu/benches/int4_prefill_route_ab.rs index 20ffd38625..fa85ecf62e 100644 --- a/crates/onnx-runtime-ep-cpu/benches/int4_prefill_route_ab.rs +++ b/crates/onnx-runtime-ep-cpu/benches/int4_prefill_route_ab.rs @@ -96,6 +96,9 @@ fn build_kernel( } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + // `PROBE_BLOCK` exists because `int4_prefill_gebp_min_rows` returns a // *different* threshold for weights whose block size the column-blocked // kernels cannot take (`INT4_PREFILL_GEBP_MIN_ROWS_UNBLOCKED`), and with diff --git a/crates/onnx-runtime-ep-cpu/benches/int8_prefill_route_ab.rs b/crates/onnx-runtime-ep-cpu/benches/int8_prefill_route_ab.rs index 3b04ed8e77..e0851b9b0b 100644 --- a/crates/onnx-runtime-ep-cpu/benches/int8_prefill_route_ab.rs +++ b/crates/onnx-runtime-ep-cpu/benches/int8_prefill_route_ab.rs @@ -130,6 +130,9 @@ fn oracle_row(activations: &[f32], weight: &Weight<'_>, row: usize, col: usize) } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + let block_size = 32usize; let shapes: Vec<(usize, usize)> = match std::env::var("PROBE_SHAPE").as_deref() { Ok("big") => vec![(4096, 11008)], diff --git a/crates/onnx-runtime-ep-cpu/benches/kernels.rs b/crates/onnx-runtime-ep-cpu/benches/kernels.rs index 961f26a4e6..157019e6a8 100644 --- a/crates/onnx-runtime-ep-cpu/benches/kernels.rs +++ b/crates/onnx-runtime-ep-cpu/benches/kernels.rs @@ -19,6 +19,10 @@ fn thread_pool(threads: usize) -> ThreadPool { } fn bench_add(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group = c.benchmark_group("add"); for (size, shape) in [ ("small", vec![1_024]), @@ -53,6 +57,10 @@ fn bench_add(c: &mut Criterion) { } fn bench_reduce_mean(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group = c.benchmark_group("reduce_mean"); for (size, shape) in [ ("small", vec![32, 128]), @@ -90,6 +98,10 @@ fn bench_reduce_mean(c: &mut Criterion) { } fn bench_gather(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group = c.benchmark_group("gather"); for (size, rows, columns, index_count) in [ ("small", 4_096, 128, 32), @@ -156,6 +168,10 @@ fn with_gemm_backend(backend: &str, f: impl FnOnce() -> T) -> T { } fn bench_matmul(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group = c.benchmark_group("matmul"); let backends = &[ "generic", @@ -267,6 +283,10 @@ fn dequantize_mxfp4_kn(n: usize, k: usize, packed: &[u8]) -> Vec { } fn bench_block_quantized_matmul_cache(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group = c.benchmark_group("block_quantized_matmul_cached_dense"); group.sample_size(15); let (m, k, n) = (1usize, 1_024usize, 1_024usize); @@ -368,6 +388,10 @@ fn block_quantized_moe_kernel(top_k: usize) -> Box { } fn bench_block_quantized_moe_cache(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group = c.benchmark_group("block_quantized_moe_cached_dense"); group.sample_size(15); let (rows, hidden, inter, experts, top_k) = (1usize, 256usize, 256usize, 4usize, 1usize); diff --git a/crates/onnx-runtime-ep-cpu/benches/matmul_nbits_prefill_ab.rs b/crates/onnx-runtime-ep-cpu/benches/matmul_nbits_prefill_ab.rs index c1f57bdebe..81536f08a3 100644 --- a/crates/onnx-runtime-ep-cpu/benches/matmul_nbits_prefill_ab.rs +++ b/crates/onnx-runtime-ep-cpu/benches/matmul_nbits_prefill_ab.rs @@ -132,6 +132,9 @@ fn median(mut samples: Vec) -> f64 { } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + let cases = [ Case { label: "int8 dense-fallback", diff --git a/crates/onnx-runtime-ep-cpu/benches/native_vs_mlas.rs b/crates/onnx-runtime-ep-cpu/benches/native_vs_mlas.rs index 558bf7177b..05e58f9540 100644 --- a/crates/onnx-runtime-ep-cpu/benches/native_vs_mlas.rs +++ b/crates/onnx-runtime-ep-cpu/benches/native_vs_mlas.rs @@ -53,6 +53,8 @@ //! regresses every concurrent session on the box. A large `ratio` with a small //! `cpu_ratio` means exactly that, and is a reason to look before graduating. +mod common; + use std::hint::black_box; use std::sync::OnceLock; use std::time::Instant; @@ -410,6 +412,9 @@ fn bench_transcendentals() { } fn main() { + // Match the decode thread topology a served session runs in (#1749). + common::init_decode_topology(); + println!("# native vs MLAS, one binary, interleaved, median of {REPS} reps"); println!( "# mlas_linked={} graduation_ratio={GRADUATION_RATIO} threads={}", diff --git a/crates/onnx-runtime-ep-cpu/benches/sdpa_simd.rs b/crates/onnx-runtime-ep-cpu/benches/sdpa_simd.rs index b6b745443c..ac07eecff9 100644 --- a/crates/onnx-runtime-ep-cpu/benches/sdpa_simd.rs +++ b/crates/onnx-runtime-ep-cpu/benches/sdpa_simd.rs @@ -11,6 +11,8 @@ //! Shapes are real encoder/decoder geometries rather than round numbers: //! Whisper-base encoder, BERT-base, ViT-L/14 and a 32/8/128 GQA decode step. +mod common; + use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main}; use onnx_runtime_ep_cpu::kernels::sdpa::{ NoBias, NoMask, ScaleMode, SdpaConfig, SdpaTensors, sdpa_f32, sdpa_f32_scalar, @@ -73,6 +75,10 @@ fn values(n: usize, seed: u64) -> Vec { } fn bench_sdpa_simd(c: &mut Criterion) { + // Match the decode thread topology a served session runs in (#1749). + // Idempotent, and called from every group member rather than just the + // first so coverage does not silently ride on `criterion_group!` order. + common::init_decode_topology(); let mut group = c.benchmark_group("sdpa_f32"); for geo in &GEOMETRIES { let batch = 1usize;