[justinchuby]

Motivation

@justinchuby noticed the CPU-EP TopKKernel isn't vectorized and does a full sort of every lane's entire axis (O(width·log width)) then truncates to k — even when k << width. Meanwhile the sampler (logits.rs::top_k_threshold) already uses an efficient bounded-heap / select_nth_unstable_by (introselect) approach. This aligns the kernel with that faster technique.

Change

Replace the per-lane full sort with:

1. select_nth_unstable_by(k-1, …) — introselect partition, O(width) average, moves the k winners to the front.
2. sort_by on just the k winners — O(k·log k).

For a 152k-wide vocab axis with k=40 this is ~an order of magnitude fewer comparisons. topk_order is a total order (ties broken by index), so the partition is deterministic; output layout (#774 k-major), index tie-breaking, and both largest/smallest orders are unchanged.

Tests

• New topk_partial_select_breaks_ties_by_index (k<width with duplicate 5.0s → must pick lowest indices [0,1,2]).
• New topk_smallest_partial_select (largest=false, ascending partition).
• All existing selection tests unchanged. cargo test -p onnx-runtime-ep-cpu --lib selection: 24 passed / 0 failed. fmt + clippy clean.

Generality: no model-name/shape special-casing. Builds on #774.

[codecov]

Codecov Report

❌ Patch coverage is 94.87179% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.22%. Comparing base (be40658) to head (d8c9127).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
...rates/onnx-runtime-ep-cpu/src/kernels/selection.rs	94.87%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #775      +/-   ##
==========================================
+ Coverage   79.89%   80.22%   +0.33%     
==========================================
  Files         331      331              
  Lines      137491   137527      +36     
  Branches   137491   137527      +36     
==========================================
+ Hits       109848   110335     +487     
+ Misses      23299    22848     -451     
  Partials     4344     4344              

Flag	Coverage Δ
cli-ort-linux	83.41% <ø> (ø)
cli-ort-windows	83.00% <ø> (+0.09%)	⬆️
mlas	81.33% <ø> (+0.04%)	⬆️
offline	80.08% <94.87%> (+0.35%)	⬆️

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

Files with missing lines	Coverage Δ
...rates/onnx-runtime-ep-cpu/src/kernels/selection.rs	82.53% <94.87%> (+0.51%)	⬆️

... and 4 files with indirect coverage changes

🚀 New features to boost your workflow:

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

[github-actions]

🔴 Benchmark Regression Detected

Comparison of criterion micro-benchmarks: PR head vs merge-base, measured on the same runner in the same job (base first → PR second).

ℹ️ Absolute times are informational only — they vary with runner load. The % change column is the reliable signal because both sides ran under identical conditions.

Status	Scenario	Base	PR	Change
🔴	gather/large_f32_threads=1-internal/131072	27.06 µs	43.66 µs	+61.4%
🔴	gather/large_f16_threads=1-internal/131072	11.82 µs	18.42 µs	+55.9%
🔴	matmul/medium_generic_f16_threads=8/32x512x512	33.73 µs	48.65 µs	+44.2%
✅	matmul/small_generic_f32_threads=1/1x256x256	37.89 µs	42.81 µs	+13.0%
✅	qwen3_sampling_processors/top_k_partial_selection	143.33 µs	157.87 µs	+10.1%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.17 ms	2.28 ms	+5.4%
✅	gather/large_bf16_threads=1-internal/131072	15.20 µs	15.76 µs	+3.7%
✅	matmul/small_generic_f32_threads=8/1x256x256	52.21 µs	53.94 µs	+3.3%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	584.82 µs	604.11 µs	+3.3%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	562.70 µs	561.45 µs	-0.2%
✅	sampling_latency/greedy_per_token	3.46 µs	3.45 µs	-0.2%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	85.98 µs	84.83 µs	-1.3%
✅	reduce_mean/small_f32_threads=1-internal/4096	17.17 µs	16.92 µs	-1.4%
✅	matmul/medium_generic_f32_threads=8/32x512x512	1.52 ms	1.50 ms	-1.6%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	100.64 µs	97.65 µs	-3.0%
✅	add/large_bf16_threads=1-internal/4194304	48.85 ms	47.23 ms	-3.3%
✅	matmul/large_generic_f32_threads=8/32x1024x1024	5.24 ms	5.06 ms	-3.4%
✅	logit_processing/seven_processor_chain_per_step	338.22 µs	323.09 µs	-4.5%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	574.86 µs	545.51 µs	-5.1%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.44 ms	9.84 ms	-5.8%
✅	matmul/medium_generic_f16_threads=1/32x512x512	38.53 µs	36.30 µs	-5.8%
✅	sampling_latency/min_p_per_token	233.91 µs	218.63 µs	-6.5%
✅	kv_cache/alloc_dealloc_pages	41.62 µs	38.68 µs	-7.1%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	4.09 ms	3.74 ms	-8.7%
✅	gather/medium_bf16_threads=1-internal/32768	2.78 µs	2.52 µs	-9.3%
✅	qwen3_sampling_processors/top_k_top_p_fast	734.28 µs	658.90 µs	-10.3%
✅	gather/medium_f16_threads=1-internal/32768	2.77 µs	2.47 µs	-11.0%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	521.59 µs	463.74 µs	-11.1%
✅	matmul/small_generic_f16_threads=8/1x256x256	44.38 µs	39.37 µs	-11.3%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	89.50 µs	79.26 µs	-11.4%
✅	add/small_f32_threads=1-internal/1024	238.4 ns	210.5 ns	-11.7%
✅	matmul/small_generic_bf16_threads=8/1x256x256	48.26 µs	42.35 µs	-12.3%
✅	add/large_f16_threads=1-internal/4194304	51.33 ms	44.95 ms	-12.4%
✅	gather/small_bf16_threads=1-internal/4096	556.7 ns	484.2 ns	-13.0%
✅	matmul/medium_generic_f32_threads=1/32x512x512	2.80 ms	2.43 ms	-13.2%
✅	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	6.64 ms	5.76 ms	-13.4%
✅	gather/medium_f32_threads=1-internal/32768	4.64 µs	4.00 µs	-14.0%
✅	tokenization/encode_tokens_per_second	448.25 µs	383.53 µs	-14.4%
✅	add/small_bf16_threads=1-internal/1024	15.68 µs	13.35 µs	-14.8%
🟢	sampling_latency/top_p_per_token	471.83 µs	400.11 µs	-15.2%
🟢	reduce_mean/large_f32_threads=1-internal/262144	1.32 ms	1.12 ms	-15.3%
🟢	add/medium_f32_threads=1-internal/262144	3.29 ms	2.78 ms	-15.5%
🟢	gather/small_f32_threads=1-internal/4096	828.6 ns	699.3 ns	-15.6%
🟢	gather/small_f16_threads=1-internal/4096	587.6 ns	494.9 ns	-15.8%
🟢	matmul/small_generic_f16_threads=1/1x256x256	44.63 µs	37.04 µs	-17.0%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	921.01 µs	762.01 µs	-17.3%
🟢	grammar_masking/llguidance_compute_mask/32	93.60 µs	76.96 µs	-17.8%
🟢	matmul/large_generic_bf16_threads=1/32x1024x1024	2.49 ms	2.02 ms	-18.9%
🟢	add/large_f32_threads=1-internal/4194304	55.72 ms	45.12 ms	-19.0%
🟢	reduce_mean/medium_f32_threads=1-internal/65536	312.37 µs	252.73 µs	-19.1%
🟢	sampling_latency/top_k_per_token	68.59 µs	54.74 µs	-20.2%
🟢	tokenization/decode_tokens_per_second	8.62 ms	6.73 ms	-21.9%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.01 ms	1.55 ms	-23.2%
🟢	add/small_f16_threads=1-internal/1024	18.26 µs	13.98 µs	-23.4%
🟢	add/medium_f16_threads=1-internal/262144	3.88 ms	2.84 ms	-26.9%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	93.73 µs	67.50 µs	-28.0%
🟢	matmul/small_generic_bf16_threads=1/1x256x256	53.74 µs	38.02 µs	-29.3%
🟢	add/medium_bf16_threads=1-internal/262144	4.40 ms	2.80 ms	-36.3%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	119.11 µs	57.91 µs	-51.4%

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.70 3.37 6.06 }

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)