[justinchuby]

Problem

CPU-EP TopKKernel built its values/indices outputs by sequential push in loop order outer → inner → rank, i.e. memory layout [outer][inner][k]. ONNX TopK output keeps the input shape with shape[axis]=k → correct layout is k-major [outer][k][inner]. These match only when inner==1 (final axis), so non-final-axis TopK with k>1 emitted wrong values+indices ordering vs the CUDA EP / ONNX Runtime. Dtype-independent, pre-existing, latent (real router TopK is final-axis).

Fix

Write each rank-r winner to its strided destination (outer*k + r)*inner + i instead of pushing sequentially. Removed the now-unused sorted field (kernel always emits sorted winners — valid for both sorted=0 and sorted=1).

Tests

• New: topk_non_final_axis_uses_k_major_layout ([2,3,2], axis=1, k=2), topk_first_axis_is_k_major ([3,2], axis=0, k=2) — assert k-major values AND indices. Old code emits e.g. [5,3,4,2] vs correct [5,4,3,2].
• Final-axis path unchanged (existing tests pass verbatim, inner==1).
• cargo test -p onnx-runtime-ep-cpu --lib selection: 22 passed / 0 failed. fmt + clippy clean.

Generality: no model-name/shape special-casing.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.36%. Comparing base (6b1e02e) to head (dc8a078).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #774      +/-   ##
==========================================
- Coverage   80.05%   79.36%   -0.70%     
==========================================
  Files         331      331              
  Lines      137462   137491      +29     
  Branches   137462   137491      +29     
==========================================
- Hits       110050   109120     -930     
- Misses      23069    24031     +962     
+ Partials     4343     4340       -3     

Flag	Coverage Δ
cli-ort-linux	83.41% <ø> (ø)
cli-ort-windows	82.90% <ø> (-0.10%)	⬇️
mlas	80.73% <ø> (ø)
offline	79.18% <100.00%> (-0.74%)	⬇️

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.02% <100.00%> (+0.64%)	⬆️

... and 6 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
🔴	kv_cache/alloc_dealloc_pages	41.00 µs	59.10 µs	+44.1%
🔴	grammar_masking/llguidance_compute_mask/32	84.83 µs	118.57 µs	+39.8%
⚠️	qwen3_sampling_processors/top_k_top_p_full_sort_baseline	5.95 ms	7.62 ms	+28.1%
⚠️	qwen3_sampling_processors/top_k_top_p_fast	670.18 µs	846.27 µs	+26.3%
⚠️	matmul/medium_generic_f16_threads=1/32x512x512	37.88 µs	45.94 µs	+21.3%
✅	qwen3_sampling_processors/top_p_fast_after_top_k	560.39 µs	633.00 µs	+13.0%
✅	logit_processing/seven_processor_chain_per_step	356.84 µs	397.88 µs	+11.5%
✅	matmul/large_generic_f16_threads=8/32x1024x1024	97.55 µs	108.32 µs	+11.0%
✅	matmul/large_generic_f32_threads=1/32x1024x1024	10.26 ms	11.31 ms	+10.2%
✅	reduce_mean/small_f32_threads=1-internal/4096	15.76 µs	17.31 µs	+9.8%
✅	block_quantized_moe_cached_dense/mxfp4_uncached_expert_dequant_each_call/rows=1,H=256,I=256,E=4,top_k=1	478.45 µs	518.72 µs	+8.4%
✅	matmul/small_generic_f32_threads=8/1x256x256	37.52 µs	40.65 µs	+8.3%
✅	gather/medium_f32_threads=1-internal/32768	3.96 µs	4.28 µs	+8.0%
✅	matmul/medium_generic_bf16_threads=8/32x512x512	586.82 µs	630.29 µs	+7.4%
✅	block_quantized_moe_cached_dense/mxfp4_cached_dense_expert_repeated_call/rows=1,H=256,I=256,E=4,top_k=1	82.13 µs	87.09 µs	+6.0%
✅	matmul/large_generic_bf16_threads=1/32x1024x1024	2.15 ms	2.26 ms	+5.3%
✅	qwen3_sampling_processors/top_k_full_sort_baseline	2.35 ms	2.45 ms	+4.3%
✅	gather/large_bf16_threads=1-internal/131072	12.53 µs	13.05 µs	+4.1%
✅	qwen3_sampling_processors/top_p_full_sort_after_top_k_baseline	3.93 ms	4.09 ms	+3.9%
✅	matmul/small_generic_f16_threads=1/1x256x256	33.08 µs	34.36 µs	+3.9%
✅	matmul/small_generic_f16_threads=8/1x256x256	33.22 µs	34.34 µs	+3.4%
✅	add/large_bf16_threads=1-internal/4194304	49.50 ms	50.81 ms	+2.6%
✅	matmul/medium_generic_f32_threads=1/32x512x512	3.07 ms	3.15 ms	+2.5%
✅	matmul/medium_generic_f32_threads=8/32x512x512	2.08 ms	2.13 ms	+2.4%
✅	gather/small_bf16_threads=1-internal/4096	498.0 ns	507.2 ns	+1.8%
✅	reduce_mean/medium_f32_threads=1-internal/65536	273.16 µs	277.81 µs	+1.7%
✅	reduce_mean/large_f32_threads=1-internal/262144	1.18 ms	1.20 ms	+1.2%
✅	add/small_bf16_threads=1-internal/1024	13.96 µs	14.09 µs	+0.9%
✅	sampling_latency/min_p_per_token	219.08 µs	220.40 µs	+0.6%
✅	gather/medium_f16_threads=1-internal/32768	2.63 µs	2.65 µs	+0.6%
✅	qwen3_sampling_processors/top_k_partial_selection	158.49 µs	159.34 µs	+0.5%
✅	sampling_latency/top_p_per_token	423.77 µs	421.67 µs	-0.5%
✅	add/large_f32_threads=1-internal/4194304	42.87 ms	42.54 ms	-0.8%
✅	matmul/small_generic_f32_threads=1/1x256x256	40.58 µs	40.06 µs	-1.3%
✅	sampling_latency/greedy_per_token	4.32 µs	4.25 µs	-1.7%
✅	gather/medium_bf16_threads=1-internal/32768	2.68 µs	2.62 µs	-2.1%
✅	add/large_f16_threads=1-internal/4194304	45.64 ms	44.43 ms	-2.7%
✅	tokenization/encode_tokens_per_second	447.84 µs	434.31 µs	-3.0%
✅	add/medium_bf16_threads=1-internal/262144	2.93 ms	2.83 ms	-3.6%
✅	matmul/small_generic_bf16_threads=8/1x256x256	47.52 µs	45.68 µs	-3.9%
✅	matmul/medium_generic_bf16_threads=1/32x512x512	632.16 µs	602.19 µs	-4.7%
✅	gather/large_f16_threads=1-internal/131072	15.07 µs	14.10 µs	-6.4%
✅	gather/small_f16_threads=1-internal/4096	525.4 ns	491.4 ns	-6.5%
✅	matmul/large_generic_f16_threads=1/32x1024x1024	93.82 µs	86.25 µs	-8.1%
✅	add/medium_f32_threads=1-internal/262144	2.96 ms	2.67 ms	-9.6%
✅	gather/small_f32_threads=1-internal/4096	900.5 ns	813.0 ns	-9.7%
✅	add/small_f16_threads=1-internal/1024	15.76 µs	14.12 µs	-10.4%
✅	matmul/medium_generic_f16_threads=8/32x512x512	44.79 µs	39.65 µs	-11.5%
✅	matmul/small_generic_bf16_threads=1/1x256x256	39.73 µs	34.98 µs	-12.0%
✅	gather/large_f32_threads=1-internal/131072	44.78 µs	39.31 µs	-12.2%
🟢	block_quantized_matmul_cached_dense/mxfp4_uncached_dequant_each_call/1x1024x1024	1.42 ms	1.18 ms	-16.4%
🟢	sampling_latency/top_k_per_token	70.29 µs	57.36 µs	-18.4%
🟢	matmul/large_generic_bf16_threads=8/32x1024x1024	2.41 ms	1.96 ms	-18.9%
🟢	matmul/large_generic_f32_threads=8/32x1024x1024	5.29 ms	4.25 ms	-19.6%
🟢	tokenization/decode_tokens_per_second	8.42 ms	6.74 ms	-20.0%
🟢	add/medium_f16_threads=1-internal/262144	3.59 ms	2.74 ms	-23.6%
🟢	add/small_f32_threads=1-internal/1024	305.4 ns	209.6 ns	-31.4%
🟢	block_quantized_matmul_cached_dense/mxfp4_cached_dense_repeated_call/1x1024x1024	117.97 µs	73.75 µs	-37.5%
🟢	block_quantized_matmul_cached_dense/mxfp4_preexpanded_dense_oncelock_like_proxy/1x1024x1024	131.03 µs	81.76 µs	-37.6%

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: { 4.63 4.05 6.01 }

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)